"""Pyodide sandbox for executing LLM-generated Prefab code.

Manages a persistent Deno subprocess running Pyodide with the full
Prefab library pre-loaded. Context managers, `.rx`, and Pydantic
validation work identically to native Python. The runner revokes Deno
filesystem and network permissions after loading its runtime and sources.
The Python interpreter and imported modules remain shared between calls.
"""

from __future__ import annotations

import asyncio
import json
import os
import shutil
import subprocess
from pathlib import Path
from typing import Any
from uuid import uuid4

RUNNER_JS = Path(__file__).parent / "runner.js"
PREFAB_SRC = Path(__file__).parent.parent  # src/prefab_ui/
DENO_LOCK = Path(__file__).parent / "deno.lock"


class PyodideSandbox:
    """Pyodide sandbox backed by a warm Deno subprocess.

    Use as an async context manager for explicit lifecycle control:

    ```python
    async with PyodideSandbox() as sandbox:
        result = await sandbox.run(code, data={"key": "value"})
    ```

    Or create a long-lived instance that starts lazily on first use:

    ```python
    sandbox = PyodideSandbox()
    result = await sandbox.run(code)  # cold start happens here
    result = await sandbox.run(code)  # warm, ~1ms
    ```

    The Deno process auto-restarts if it crashes between calls.
    """

    def __init__(self) -> None:
        self._process: subprocess.Popen[bytes] | None = None
        self._lock = asyncio.Lock()

    async def __aenter__(self) -> PyodideSandbox:
        await self._start()
        return self

    async def __aexit__(self, *args: Any) -> None:
        self._stop()

    async def _start(self) -> None:
        deno = shutil.which("deno")
        if deno is None:
            raise RuntimeError(
                "Deno is required for the Prefab sandbox. "
                "Install it from https://deno.land"
            )

        loop = asyncio.get_event_loop()
        env = {
            **os.environ,
            "DENO_V8_FLAGS": "--max-old-space-size=4096",
            "DENO_UNSTABLE_DETECT_CJS": "1",
        }
        starting = loop.run_in_executor(
            None,
            lambda: subprocess.Popen(
                [
                    deno,
                    "run",
                    "--no-config",
                    f"--lock={DENO_LOCK}",
                    "--frozen",
                    "--unstable-detect-cjs",
                    "--no-prompt",
                    "--allow-read",
                    "--allow-write",
                    "--allow-net=cdn.jsdelivr.net:443",
                    str(RUNNER_JS),
                    str(PREFAB_SRC),
                ],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                env=env,
            ),
        )

        try:
            proc = await asyncio.shield(starting)
        except asyncio.CancelledError:
            # Popen may still be running in its executor thread. Keep its
            # result available so the process can be reclaimed when it returns.
            def discard_startup(
                future: asyncio.Future[subprocess.Popen[bytes]],
            ) -> None:
                if not future.cancelled() and future.exception() is None:
                    self._stop_process(future.result())

            starting.add_done_callback(discard_startup)
            raise

        try:
            # Wait for the runner to signal readiness, collecting stderr
            assert proc.stderr is not None
            proc_stderr = proc.stderr
            stderr_lines: list[str] = []
            while True:
                line = await loop.run_in_executor(None, proc_stderr.readline)
                if not line and proc.poll() is not None:
                    err = "\n".join(stderr_lines)
                    if len(err) > 2000:
                        err = f"{err[:1000]}\n...\n{err[-1000:]}"
                    raise RuntimeError(f"Pyodide sandbox failed to start:\n{err}")
                decoded = line.decode().rstrip()
                stderr_lines.append(decoded)
                if "pyodide:ready" in decoded:
                    break
        except (Exception, asyncio.CancelledError):
            self._stop_process(proc)
            raise

        self._process = proc

    def _stop(self) -> None:
        if self._process is not None:
            proc = self._process
            self._process = None
            self._stop_process(proc)

    @staticmethod
    def _stop_process(proc: subprocess.Popen[bytes]) -> None:
        if proc.poll() is None:
            proc.terminate()
        try:
            proc.communicate(timeout=5)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.communicate()
        finally:
            for stream in (proc.stdin, proc.stdout, proc.stderr):
                if stream is not None:
                    stream.close()

    async def run(
        self,
        code: str,
        *,
        data: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Execute Prefab code and return the wire protocol JSON.

        Args:
            code: Python code that builds a Prefab component tree.
                Must assign a `PrefabApp` or `Component` to a variable.
            data: Values injected as variables in the sandbox namespace.

        Returns:
            Prefab wire protocol dict (`$prefab`, `view`, `state`, etc.)

        Cancellation during startup or execution discards the worker. The next
        call starts a fresh process. Use `asyncio.wait_for()` to apply a timeout.

        Raises:
            RuntimeError: If the code fails or the sandbox is unavailable.
        """
        async with self._lock:
            # Lazy start / auto-restart
            if self._process is None or self._process.poll() is not None:
                await self._start()

            proc = self._process
            if proc is None or proc.stdin is None or proc.stdout is None:
                raise RuntimeError("Sandbox failed to start")

            request_id = uuid4().hex
            request = (
                json.dumps({"id": request_id, "code": code, "data": data or {}}) + "\n"
            )
            loop = asyncio.get_event_loop()
            try:
                proc.stdin.write(request.encode())
                proc.stdin.flush()
                line = await loop.run_in_executor(None, proc.stdout.readline)
            except (OSError, asyncio.CancelledError):
                self._stop()
                raise

            if not line:
                self._stop()
                raise RuntimeError("Sandbox process died unexpectedly")

            try:
                response = json.loads(line)
            except (ValueError, UnicodeDecodeError) as exc:
                self._stop()
                raise RuntimeError("Invalid sandbox response") from exc

            if (
                not isinstance(response, dict)
                or response.get("id") != request_id
                or not (
                    isinstance(response.get("result"), dict)
                    or isinstance(response.get("error"), str)
                )
            ):
                self._stop()
                raise RuntimeError("Invalid sandbox response")
            if "error" in response:
                raise RuntimeError(response["error"])
            return response["result"]