"""Jupyter kernel manager for executing Python code with persistent state."""

import shutil
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path

from jupyter_client import KernelManager as JupyterKernelManager
from jupyter_client import BlockingKernelClient

from .config import get_kernel_venv_dir, get_kernel_connection_file

DEFAULT_PACKAGES = [
    "ipykernel",
    "jupyter_client",
    "polars",
    "pandas",
    "numpy",
    "matplotlib",
    "seaborn",
    "pyyaml",
    "python-dotenv",
]


@dataclass
class ExecutionResult:
    """Result of code execution in the kernel."""

    success: bool
    output: str
    error: str | None = None


class KernelManager:
    """Manages a Jupyter kernel for Python code execution."""

    def __init__(
        self,
        venv_dir: Path | None = None,
        kernel_name: str = "astro-ai-kernel",
        packages: list[str] | None = None,
    ):
        self.venv_dir = venv_dir or get_kernel_venv_dir()
        self.kernel_name = kernel_name
        self.packages = packages or DEFAULT_PACKAGES.copy()
        self.connection_file = get_kernel_connection_file()
        self._km: JupyterKernelManager | None = None

    @property
    def python_path(self) -> Path:
        if sys.platform == "win32":
            return self.venv_dir / "Scripts" / "python.exe"
        return self.venv_dir / "bin" / "python"

    @property
    def is_running(self) -> bool:
        if not self.connection_file.exists():
            return False
        try:
            kc = BlockingKernelClient()
            kc.load_connection_file(str(self.connection_file))
            kc.start_channels()
            try:
                kc.wait_for_ready(timeout=2)
                return True
            except Exception:
                return False
            finally:
                kc.stop_channels()
        except Exception:
            return False

    def ensure_environment(self, extra_packages: list[str] | None = None) -> None:
        if not shutil.which("uv"):
            raise RuntimeError(
                "uv is not installed.\n"
                "Install with: curl -LsSf https://astral.sh/uv/install.sh | sh"
            )

        packages = self.packages.copy()
        if extra_packages:
            packages.extend(extra_packages)

        if not self.venv_dir.exists():
            print(f"Creating environment at {self.venv_dir}")
            subprocess.run(
                ["uv", "venv", str(self.venv_dir), "--seed"],
                check=True,
                capture_output=True,
            )

        print("Installing packages...")
        subprocess.run(
            ["uv", "pip", "install", "--python", str(self.python_path)] + packages,
            check=True,
            capture_output=True,
        )

        # Register kernel
        try:
            subprocess.run(
                [
                    str(self.python_path),
                    "-m",
                    "ipykernel",
                    "install",
                    "--user",
                    "--name",
                    self.kernel_name,
                    "--display-name",
                    "Data Analysis Kernel",
                ],
                capture_output=True,
                timeout=30,
            )
        except Exception:
            pass

    def start(
        self,
        env_vars: dict[str, str] | None = None,
        extra_packages: list[str] | None = None,
    ) -> None:
        if self.is_running:
            print("Kernel already running")
            return

        self.ensure_environment(extra_packages=extra_packages)
        print("Starting kernel...")

        self._km = JupyterKernelManager(kernel_name=self.kernel_name)

        if env_vars:
            import os

            for key, value in env_vars.items():
                os.environ[key] = value

        self._km.start_kernel(extra_arguments=["--IPKernelApp.parent_handle=0"])

        self.connection_file.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy(self._km.connection_file, self.connection_file)

        kc = self._km.client()
        kc.start_channels()
        try:
            kc.wait_for_ready(timeout=10)
        except Exception as e:
            self.stop()
            raise RuntimeError(f"Kernel failed: {e}") from e
        finally:
            kc.stop_channels()

        self._km = None
        print(f"Kernel started ({self.connection_file})")

    def stop(self) -> None:
        if not self.connection_file.exists():
            print("Kernel not running")
            return

        try:
            kc = BlockingKernelClient()
            kc.load_connection_file(str(self.connection_file))
            kc.start_channels()
            kc.shutdown()
            kc.stop_channels()
        except Exception:
            pass

        if self.connection_file.exists():
            self.connection_file.unlink()
        print('{"message": "Kernel stopped"}')

    def execute(self, code: str, timeout: float = 30.0) -> ExecutionResult:
        if not self.connection_file.exists():
            return ExecutionResult(
                False, "", "Kernel not running. Start with: uv run scripts/cli.py start"
            )

        kc = BlockingKernelClient()
        kc.load_connection_file(str(self.connection_file))
        kc.start_channels()

        try:
            kc.wait_for_ready(timeout=5)
        except Exception as e:
            kc.stop_channels()
            return ExecutionResult(False, "", f"Kernel not responding: {e}")

        msg_id = kc.execute(code, silent=False, store_history=True)

        output_parts: list[str] = []
        error_msg: str | None = None
        status = "ok"
        deadline = time.time() + timeout
        done = False

        while time.time() < deadline and not done:
            try:
                msg = kc.get_iopub_msg(timeout=min(1.0, deadline - time.time()))
                if msg["parent_header"].get("msg_id") != msg_id:
                    continue

                msg_type = msg["msg_type"]
                content = msg["content"]

                if msg_type == "stream":
                    output_parts.append(content["text"])
                elif msg_type == "execute_result":
                    output_parts.append(content["data"].get("text/plain", ""))
                elif msg_type == "error":
                    error_msg = "\n".join(content["traceback"])
                    status = "error"
                elif msg_type == "status" and content["execution_state"] == "idle":
                    done = True
            except Exception:
                continue

        kc.stop_channels()

        if not done:
            return ExecutionResult(
                False, "".join(output_parts), f"Timeout after {timeout}s"
            )

        return ExecutionResult(status == "ok", "".join(output_parts), error_msg)

    def status(self) -> dict:
        info = {
            "running": False,
            "connection_file": str(self.connection_file),
            "responsive": False,
        }
        if not self.connection_file.exists():
            return info
        info["running"] = True
        try:
            kc = BlockingKernelClient()
            kc.load_connection_file(str(self.connection_file))
            kc.start_channels()
            try:
                kc.wait_for_ready(timeout=2)
                info["responsive"] = True
            except Exception:
                pass
            finally:
                kc.stop_channels()
        except Exception:
            pass
        return info

    def install_packages(self, packages: list[str]) -> tuple[bool, str]:
        """Install additional packages into the kernel environment.

        Args:
            packages: List of package specs (e.g., ['plotly>=5.0', 'scipy'])

        Returns:
            Tuple of (success, message)
        """
        if not packages:
            return False, "No packages specified"

        if not shutil.which("uv"):
            return False, "uv is not installed"

        try:
            result = subprocess.run(
                ["uv", "pip", "install", "--python", str(self.python_path)] + packages,
                capture_output=True,
                text=True,
            )
            if result.returncode == 0:
                return True, f"Installed: {', '.join(packages)}"
            else:
                return False, f"Failed: {result.stderr}"
        except Exception as e:
            return False, f"Error: {e}"
