#!/usr/bin/env python3
"""Whalent-owned Jupyter kernel sidecar.

Stdin/stdout are NDJSON. Stdout is protocol-only; diagnostics go to stderr.
The process is launched from the selected Python environment so imports and the
kernel executable are guaranteed to belong to that environment.
"""

from __future__ import annotations

import json
import os
import queue
import signal
import sys
import threading
import time
import traceback
from typing import Any

from jupyter_client import KernelManager
from jupyter_client.kernelspec import KernelSpec


WRITE_LOCK = threading.Lock()
KERNEL_LOCK = threading.RLock()
KM: KernelManager | None = None
KC: Any = None
GENERATION = 0
STATE = "dead"
CURRENT_CELL: str | None = None
SHUTTING_DOWN = False
MAX_FRAME_BYTES = 8 * 1024 * 1024


def emit(payload: dict[str, Any]) -> None:
    encoded = json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
    if len(encoded.encode("utf-8")) > MAX_FRAME_BYTES:
        encoded = json.dumps(
            {"type": "protocol.error", "error": "output frame exceeds 8 MiB"},
            ensure_ascii=False,
            separators=(",", ":"),
        )
    with WRITE_LOCK:
        sys.stdout.write(encoded + "\n")
        sys.stdout.flush()


def reply(request_id: str, ok: bool, result: Any = None, error: str = "") -> None:
    payload: dict[str, Any] = {"type": "rpc.reply", "request_id": request_id, "ok": ok}
    if ok:
        payload["result"] = result
    else:
        payload["error"] = {"message": error or "sidecar request failed"}
    emit(payload)


def safe_status() -> dict[str, Any]:
    pid = None
    with KERNEL_LOCK:
        if KM is not None:
            provisioner = getattr(KM, "provisioner", None)
            pid = getattr(provisioner, "pid", None)
    return {
        "state": STATE,
        "generation": GENERATION,
        "current_cell": CURRENT_CELL,
        "pid": int(pid) if isinstance(pid, int) else None,
    }


def start_kernel() -> dict[str, Any]:
    global KM, KC, GENERATION, STATE
    with KERNEL_LOCK:
        if KM is not None and KM.is_alive():
            return safe_status()
        STATE = "starting"
        manager = KernelManager()
        # The sidecar itself runs under the selected interpreter. Pin ipykernel
        # to the same executable via an explicit KernelSpec: assigning
        # `manager.kernel_cmd` is a silent no-op on jupyter_client >= 7 (the
        # trait is gone), while `_kernel_spec` is the cache slot behind the
        # KernelManager.kernel_spec property and keeps its meaning on 5.x-8.x.
        manager._kernel_spec = KernelSpec(
            argv=[sys.executable, "-m", "ipykernel_launcher", "-f", "{connection_file}"],
            display_name="Whalent (pinned)",
            language="python",
        )
        # Kernel stdout/stderr go to the sidecar's stderr so raw fd writes from
        # user code can never corrupt the NDJSON protocol on stdout.
        manager.start_kernel(stdout=sys.stderr, stderr=sys.stderr)
        client = manager.client()
        client.start_channels()
        client.wait_for_ready(timeout=30)
        KM = manager
        KC = client
        GENERATION += 1
        STATE = "idle"
        emit({"type": "kernel.status", **safe_status()})
        return safe_status()


def output_from_message(message: dict[str, Any]) -> dict[str, Any] | None:
    msg_type = str(message.get("header", {}).get("msg_type", ""))
    content = message.get("content", {})
    if msg_type == "stream":
        return {
            "outputType": "stream",
            "name": str(content.get("name", "stdout")),
            "text": str(content.get("text", "")),
        }
    if msg_type in ("display_data", "execute_result"):
        data = content.get("data", {})
        safe_data = {
            str(key): value if isinstance(value, str) else json.dumps(value, ensure_ascii=False)
            for key, value in data.items()
            if isinstance(key, str)
        }
        transient = content.get("transient", {})
        return {
            "outputType": msg_type,
            "data": safe_data,
            "executionCount": content.get("execution_count"),
            "displayId": transient.get("display_id") if isinstance(transient, dict) else None,
        }
    if msg_type == "error":
        return {
            "outputType": "error",
            "name": str(content.get("ename", "Error")),
            "value": str(content.get("evalue", "")),
            "traceback": [str(line) for line in content.get("traceback", [])],
        }
    if msg_type == "clear_output":
        return {"outputType": "clear_output", "wait": bool(content.get("wait"))}
    return None


def execute_worker(request_id: str, cell_id: str, code: str, timeout_seconds: float) -> None:
    global STATE, CURRENT_CELL
    generation = GENERATION
    try:
        start_kernel()
        with KERNEL_LOCK:
            # Capture the generation this execute belongs to (start_kernel may
            # have revived a dead kernel and bumped it); every frame below is
            # stamped with it so the daemon can fence stale generations.
            generation = GENERATION
            client = KC
            if client is None:
                raise RuntimeError("kernel client unavailable")
            msg_id = client.execute(code, stop_on_error=True, allow_stdin=False)
            STATE = "busy"
            CURRENT_CELL = cell_id
        emit({"type": "execution.started", "request_id": request_id, "cell_id": cell_id, **safe_status(), "generation": generation})
        deadline = time.monotonic() + timeout_seconds
        execution_count = None
        shell_done = False
        idle = False
        while time.monotonic() < deadline and not (shell_done and idle):
            try:
                message = client.get_iopub_msg(timeout=0.2)
                if message.get("parent_header", {}).get("msg_id") == msg_id:
                    msg_type = message.get("header", {}).get("msg_type")
                    if msg_type == "status" and message.get("content", {}).get("execution_state") == "idle":
                        idle = True
                    output = output_from_message(message)
                    if output:
                        if output.get("outputType") == "clear_output":
                            emit({"type": "output.clear", "cell_id": cell_id, "generation": generation, "wait": output.get("wait", False)})
                        else:
                            emit({"type": "output", "cell_id": cell_id, "generation": generation, "output": output})
            except queue.Empty:
                pass
            if not shell_done:
                try:
                    shell = client.get_shell_msg(timeout=0.01)
                    if shell.get("parent_header", {}).get("msg_id") == msg_id:
                        execution_count = shell.get("content", {}).get("execution_count")
                        shell_done = True
                except queue.Empty:
                    pass
        if not (shell_done and idle):
            try:
                with KERNEL_LOCK:
                    if KM is not None:
                        KM.interrupt_kernel()
            except Exception:
                pass
            raise TimeoutError(f"cell execution exceeded {int(timeout_seconds)} seconds")
        with KERNEL_LOCK:
            if CURRENT_CELL == cell_id:
                STATE = "idle"
                CURRENT_CELL = None
        emit({
            "type": "execution.reply",
            "request_id": request_id,
            "ok": True,
            "cell_id": cell_id,
            "execution_count": execution_count,
            **safe_status(),
            "generation": generation,
        })
    except Exception as exc:
        # Only release the claim if this worker still owns it: a restart may
        # already have handed the kernel to a newer execute.
        with KERNEL_LOCK:
            if CURRENT_CELL == cell_id:
                STATE = "idle" if KM is not None and KM.is_alive() else "dead"
                CURRENT_CELL = None
        emit({
            "type": "execution.reply",
            "request_id": request_id,
            "ok": False,
            "cell_id": cell_id,
            "error": {"message": str(exc)},
            **safe_status(),
            "generation": generation,
        })


def interrupt_kernel() -> dict[str, Any]:
    global STATE
    with KERNEL_LOCK:
        if KM is None or not KM.is_alive():
            raise RuntimeError("kernel is not running")
        KM.interrupt_kernel()
    STATE = "idle"
    emit({"type": "execution.interrupted", **safe_status()})
    return safe_status()


def shutdown_kernel() -> dict[str, Any]:
    global KM, KC, STATE, CURRENT_CELL
    with KERNEL_LOCK:
        client, manager = KC, KM
        KC = None
        KM = None
    if client is not None:
        try:
            client.stop_channels()
        except Exception:
            pass
    if manager is not None:
        try:
            manager.shutdown_kernel(now=True)
        except Exception:
            try:
                manager.cleanup_resources(restart=False)
            except Exception:
                pass
    STATE = "dead"
    CURRENT_CELL = None
    emit({"type": "kernel.status", **safe_status()})
    return safe_status()


def handle(message: dict[str, Any]) -> None:
    global STATE, CURRENT_CELL
    request_id = str(message.get("id", ""))
    method = str(message.get("method", ""))
    params = message.get("params") if isinstance(message.get("params"), dict) else {}
    try:
        if method == "start":
            reply(request_id, True, start_kernel())
        elif method == "execute":
            cell_id = str(params.get("cell_id", ""))
            code = str(params.get("code", ""))
            if not cell_id:
                raise ValueError("cell_id required")
            with KERNEL_LOCK:
                if CURRENT_CELL is not None:
                    raise ValueError(f"kernel busy: cell {CURRENT_CELL} still running")
                # Claim the kernel before the worker thread exists so two
                # racing execute requests can never drain IOPub concurrently.
                STATE = "busy"
                CURRENT_CELL = cell_id
            try:
                threading.Thread(
                    target=execute_worker,
                    args=(request_id, cell_id, code, max(1.0, min(float(params.get("timeout_seconds", 1800)), 3600.0))),
                    daemon=True,
                ).start()
            except BaseException:
                with KERNEL_LOCK:
                    STATE = "idle" if KM is not None and KM.is_alive() else "dead"
                    CURRENT_CELL = None
                raise
        elif method == "interrupt":
            reply(request_id, True, interrupt_kernel())
        elif method == "restart":
            shutdown_kernel()
            reply(request_id, True, start_kernel())
        elif method == "shutdown":
            reply(request_id, True, shutdown_kernel())
        elif method == "status":
            reply(request_id, True, safe_status())
        else:
            raise ValueError(f"unsupported method: {method}")
    except Exception as exc:
        reply(request_id, False, error=str(exc))


def cleanup(*_: Any) -> None:
    global SHUTTING_DOWN
    if SHUTTING_DOWN:
        return
    SHUTTING_DOWN = True
    try:
        shutdown_kernel()
    finally:
        raise SystemExit(0)


def main() -> None:
    signal.signal(signal.SIGTERM, cleanup)
    if hasattr(signal, "SIGINT"):
        signal.signal(signal.SIGINT, cleanup)
    emit({"type": "sidecar.ready", "pid": os.getpid(), "protocol": 1})
    for line in sys.stdin:
        try:
            message = json.loads(line)
            if not isinstance(message, dict):
                raise ValueError("request must be an object")
            handle(message)
        except Exception as exc:
            print(f"[notebook-sidecar] {exc}", file=sys.stderr, flush=True)
    cleanup()


if __name__ == "__main__":
    main()