# Copyright 2026 Marimo. All rights reserved.

"""Kernel manager for marimo-lsp."""

from __future__ import annotations

import subprocess
import typing
from pathlib import Path

import marimo._ipc as ipc
from marimo._config.settings import GLOBAL_SETTINGS
from marimo._runtime.commands import AppMetadata
from marimo._session.managers import KernelManagerImpl
from marimo._session.model import SessionMode
from marimo._session.queue import ProcessLike

from marimo_lsp.loggers import get_logger

if typing.TYPE_CHECKING:
    from marimo._config.manager import MarimoConfigManager
    from marimo._ipc.types import ConnectionInfo
    from marimo._session.types import QueueManager

    from marimo_lsp.app_file_manager import LspAppFileManager


logger = get_logger()


class InvalidWorkingDirectoryError(ValueError):
    """The requested kernel working directory is unsafe or unavailable."""

    def __init__(self, path: Path, reason: str) -> None:
        message = f"Kernel working directory {reason}: {path}"
        super().__init__(message)


def _validate_working_directory(path: Path) -> None:
    if not path.is_absolute():
        raise InvalidWorkingDirectoryError(path, "must be absolute")
    if not path.exists():
        raise InvalidWorkingDirectoryError(path, "does not exist")
    if not path.is_dir():
        raise InvalidWorkingDirectoryError(path, "is not a directory")


def launch_kernel(
    executable: str,
    args: ipc.KernelArgs,
    cwd: str | None = None,
) -> Process:
    """Launch kernel as a subprocess."""
    cmd = [executable, "-m", "marimo._ipc.launch_kernel"]
    logger.info(f"Launching kernel subprocess: {' '.join(cmd)}")
    logger.debug(f"Connection info: {args.connection_info}")
    if cwd:
        logger.debug(f"Setting kernel working directory to: {cwd}")

    process = subprocess.Popen(  # noqa: S603
        cmd,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        cwd=cwd,
    )

    assert process.stdin, "Expect subprocess stdin pipe"
    assert process.stdout, "Expect subprocess stdout pipe"
    assert process.stderr, "Expect subprocess stderr pipe"

    # Send over stdin
    process.stdin.write(args.encode_json())
    process.stdin.flush()
    process.stdin.close()

    logger.debug("Waiting for KERNEL_READY signal from kernel subprocess")

    # Wait for "KERNEL_READY" message
    ready_line = process.stdout.readline().decode("utf-8").strip()
    if ready_line != "KERNEL_READY":
        exit_code = process.poll()
        stderr = process.stderr.read().decode("utf-8", errors="replace")

        if exit_code is not None and exit_code != 0:
            msg = f"Kernel failed to start (exit code {exit_code}): {stderr}"
            logger.error(msg)
            raise RuntimeError(msg)

        msg = (
            f"Invalid kernel response. Expected 'KERNEL_READY', got: '{ready_line}'. "
            f"Stderr: {stderr}"
        )
        logger.error(msg)
        process.terminate()
        raise RuntimeError(msg)

    logger.info(f"Kernel subprocess started successfully with PID: {process.pid}")
    return Process(inner=process)


class Manager(KernelManagerImpl):
    """Kernel manager for marimo-lsp."""

    def __init__(  # noqa: PLR0913
        self,
        *,
        executable: str,
        queue_manager: QueueManager,
        app_file_manager: LspAppFileManager,
        config_manager: MarimoConfigManager,
        connection_info: ConnectionInfo,
        working_directory: str,
    ) -> None:
        super().__init__(
            # NB: Leaky abstraction. Mode affects internal behavior of
            # `KernelManager` based on whether `stream_queue` or
            # `socket_addr` is used.
            #
            # We use RUN even though VS Code is conceptually more like EDIT.
            # In EDIT mode, the kernel creates a `multiprocessing.Connection`,
            # which we'd need to proxy. RUN mode uses existing `stream_queue`
            # for message distribution which aligns with our ZeroMQ
            # architecture.
            mode=SessionMode.RUN,
            config_manager=config_manager,
            configs=app_file_manager.app.cell_manager.config_map(),
            queue_manager=queue_manager,
            app_metadata=AppMetadata(
                query_params={},
                filename=app_file_manager.path,
                cli_args={},
                argv=None,
                app_config=app_file_manager.app.config,
            ),
            redirect_console_to_browser=True,
            virtual_file_storage=None,
        )
        self.executable = executable
        self.connection_info = connection_info
        self.working_directory = working_directory

    def start_kernel(self) -> None:
        """Start an instance of the marimo kernel using ZeroMQ IPC."""
        working_directory = Path(self.working_directory)
        _validate_working_directory(working_directory)
        logger.info(f"Kernel working directory: {working_directory}")
        self.kernel_task = launch_kernel(
            executable=self.executable,
            args=ipc.KernelArgs(
                connection_info=self.connection_info,
                configs=self.configs,
                app_metadata=self.app_metadata,
                user_config=self.config_manager.get_config(hide_secrets=False),
                log_level=GLOBAL_SETTINGS.LOG_LEVEL,
                profile_path=self.profile_path,
            ),
            cwd=str(working_directory),
        )


class Process(ProcessLike):
    """Wraps `subprocess.Popen` as a `ProcessLike`.

    Provides the `ProcessLike` protocol required by marimo's KernelManager.
    """

    # Timeout in seconds to wait for graceful termination before force killing
    TERMINATE_TIMEOUT = 2.0

    def __init__(self, inner: subprocess.Popen[bytes]) -> None:
        """Initialize with a subprocess.Popen instance."""
        self.inner = inner

    @property
    def pid(self) -> int | None:
        """Get the process ID."""
        return self.inner.pid

    def is_alive(self) -> bool:
        """Check if the process is still running."""
        return self.inner.poll() is None

    def terminate(self) -> None:
        """Terminate the process, force killing if it doesn't exit gracefully.

        First sends SIGTERM and waits briefly for the process to exit.
        If the process is still running after the timeout, sends SIGKILL
        to ensure cleanup (prevents orphaned kernel processes on extension
        host restart or unexpected shutdown).
        """
        pid = self.inner.pid
        if not self.is_alive():
            logger.debug(f"Kernel process {pid} already terminated")
            return

        logger.info(f"Terminating kernel process {pid}")
        self.inner.terminate()

        try:
            # Wait for graceful termination
            self.inner.wait(timeout=self.TERMINATE_TIMEOUT)
            logger.info(f"Kernel process {pid} terminated gracefully")
        except subprocess.TimeoutExpired:
            # Process didn't exit in time, force kill
            logger.warning(
                f"Kernel process {pid} did not terminate gracefully, sending SIGKILL"
            )
            self.inner.kill()
            try:
                self.inner.wait(timeout=1.0)
                logger.info(f"Kernel process {pid} killed")
            except subprocess.TimeoutExpired:
                logger.exception(f"Failed to kill kernel process {pid}")
