"""
Core-side runtime supervisor + IPC client for plugin subprocesses.
"""
from __future__ import annotations

import json
import os
import queue
import resource
import subprocess
import sys
import threading
import time
import uuid
from dataclasses import dataclass
from typing import Any, Callable

from plugin_ipc_protocol import (
    ACTION_CANCEL_REQUEST,
    ACTION_HANDSHAKE,
    ACTION_HEALTH,
    EVENT_TOOL_STREAM_TIMEOUT,
    IpcErrorCode,
    decode_frame,
    encode_frame,
    make_request,
)

# Default basenames allowed for subprocess.run in sandboxed plugins when manifest omits shell_allowlist.
_DEFAULT_SUBPROCESS_ALLOWLIST = [
    "python",
    "python3",
    "bash",
    "sh",
    "ffmpeg",
    "git",
    "yt-dlp",
]


class RuntimeState:
    DISCOVERED = "discovered"
    STARTING = "starting"
    HEALTHY = "healthy"
    DEGRADED = "degraded"
    CRASHED = "crashed"
    STOPPING = "stopping"
    STOPPED = "stopped"
    QUARANTINED = "quarantined"


@dataclass
class RuntimeMetrics:
    crash_count: int = 0
    timeout_count: int = 0
    invoke_count: int = 0
    total_invoke_ms: float = 0.0
    p95_window: list[float] | None = None

    def __post_init__(self) -> None:
        if self.p95_window is None:
            self.p95_window = []

    def add_invoke(self, ms: float) -> None:
        self.invoke_count += 1
        self.total_invoke_ms += float(ms)
        self.p95_window.append(float(ms))
        if len(self.p95_window) > 100:
            self.p95_window = self.p95_window[-100:]

    @property
    def p50_ms(self) -> float:
        vals = sorted(self.p95_window or [])
        if not vals:
            return 0.0
        return float(vals[len(vals) // 2])

    @property
    def p95_ms(self) -> float:
        vals = sorted(self.p95_window or [])
        if not vals:
            return 0.0
        idx = max(0, min(len(vals) - 1, int(len(vals) * 0.95) - 1))
        return float(vals[idx])


class PluginRuntimeHost:
    def __init__(
        self,
        *,
        plugin_slug: str,
        plugin_name: str,
        plugin_dir: str,
        plugin_hash: str,
        permissions: dict[str, Any] | None = None,
        settings_get: Callable[[str, Any], Any] | None = None,
        settings_set: Callable[[str, Any], None] | None = None,
        on_stream_event: Callable[[dict[str, Any]], None] | None = None,
        queue_limit: int = 64,
    ) -> None:
        self.plugin_slug = plugin_slug
        self.plugin_name = plugin_name
        self.plugin_dir = plugin_dir
        self.plugin_hash = plugin_hash
        self.permissions = permissions or {}
        self.settings_get = settings_get
        self.settings_set = settings_set
        self.on_stream_event = on_stream_event
        self.queue_limit = max(1, int(queue_limit))
        self._auth = uuid.uuid4().hex
        self._session_nonce = uuid.uuid4().hex
        self._seq = 0
        self._last_inbound_seq = 0
        self._proc: subprocess.Popen | None = None
        self._state = RuntimeState.DISCOVERED
        self._state_lock = threading.RLock()
        self._pending: dict[str, queue.Queue] = {}
        self._stream_handlers: dict[str, Callable[[dict[str, Any]], None]] = {}
        self._pending_lock = threading.Lock()
        self._writer_lock = threading.Lock()
        self._reader_thread: threading.Thread | None = None
        self._stderr_thread: threading.Thread | None = None
        self._stop_evt = threading.Event()
        self._inflight_sem = threading.BoundedSemaphore(self.queue_limit)
        self.last_exit_code: int | None = None
        self.last_error: str = ""
        self.last_heartbeat_ms: int | None = None
        self.quarantined_reason: str = ""
        self.last_failure_code: str = ""
        self.metrics = RuntimeMetrics()
        self._bootstrap: dict[str, Any] = {}
        self._backoff_s = 1.0
        self._consecutive_failures = 0
        self._max_failures_before_quarantine = 3
        self._lifecycle_lock = threading.RLock()

    @property
    def consecutive_failures(self) -> int:
        return int(self._consecutive_failures)

    @property
    def current_backoff_s(self) -> float:
        return float(self._backoff_s)

    @property
    def queue_depth(self) -> int:
        with self._pending_lock:
            return len(self._pending)

    @property
    def state(self) -> str:
        with self._state_lock:
            return self._state

    def _set_state(self, value: str) -> None:
        allowed = {
            RuntimeState.DISCOVERED: {RuntimeState.STARTING, RuntimeState.STOPPED, RuntimeState.QUARANTINED},
            RuntimeState.STARTING: {RuntimeState.HEALTHY, RuntimeState.DEGRADED, RuntimeState.CRASHED, RuntimeState.STOPPING},
            RuntimeState.HEALTHY: {RuntimeState.DEGRADED, RuntimeState.STOPPING, RuntimeState.QUARANTINED, RuntimeState.CRASHED},
            RuntimeState.DEGRADED: {RuntimeState.STARTING, RuntimeState.STOPPING, RuntimeState.CRASHED, RuntimeState.QUARANTINED},
            RuntimeState.CRASHED: {RuntimeState.STARTING, RuntimeState.STOPPING, RuntimeState.QUARANTINED},
            RuntimeState.STOPPING: {RuntimeState.STOPPED, RuntimeState.CRASHED},
            RuntimeState.STOPPED: {RuntimeState.STARTING, RuntimeState.QUARANTINED},
            RuntimeState.QUARANTINED: {RuntimeState.QUARANTINED},
        }
        with self._state_lock:
            prev = self._state
            if value == prev:
                return
            if value not in allowed.get(prev, set()):
                # Never coerce to DEGRADED here: that is invalid from terminal states (e.g. QUARANTINED)
                # and can undo quarantine when another thread races (e.g. reader -> CRASHED).
                self.last_error = f"invalid lifecycle transition: {prev} -> {value} (state unchanged)"
                return
            self._state = value

    def quarantine(self, reason: str) -> None:
        self.quarantined_reason = str(reason or "quarantined")
        self._set_state(RuntimeState.QUARANTINED)
        self.stop()

    def _next_seq(self) -> int:
        with self._writer_lock:
            self._seq += 1
            return self._seq

    def _build_env(self) -> dict[str, str]:
        env: dict[str, str] = {}
        # Minimal inherited env; preserve PATH for interpreter behavior.
        for key in ("PATH", "HOME", "LANG", "LC_ALL", "TMPDIR"):
            if key in os.environ:
                env[key] = os.environ[key]
        # Dashboard/bridge feeds need the same localhost auth + URL as the main backend.
        for key in (
            "MUGI_CALENDAR_URL",
            "MUGI_AUTH_TOKEN",
            "MUGI_REPO_ROOT",
        ):
            if key in os.environ:
                env[key] = os.environ[key]
        self._merge_bridge_creds_for_subprocess(env)
        env["PYTHONUNBUFFERED"] = "1"
        env["PYTHONPATH"] = os.getcwd()
        return env

    @staticmethod
    def _merge_bridge_creds_for_subprocess(env: dict[str, str]) -> None:
        """If the parent env is missing loopback bridge vars, read ~/.mugi/.macos_bridge_runner_env.json.

        Swift writes this when the app starts the backend so LaunchAgent runners match; plugin
        subprocesses can otherwise miss MUGI_AUTH_TOKEN if the parent only had the file.
        """
        if (env.get("MUGI_AUTH_TOKEN") or "").strip() and (env.get("MUGI_CALENDAR_URL") or "").strip():
            return
        path = os.path.join(os.path.expanduser("~"), ".mugi", ".macos_bridge_runner_env.json")
        if not os.path.isfile(path):
            return
        try:
            with open(path, encoding="utf-8") as f:
                data = json.load(f)
        except Exception:
            return
        if not isinstance(data, dict):
            return
        url = str(data.get("MUGI_CALENDAR_URL") or "").strip().rstrip("/")
        token = str(data.get("MUGI_AUTH_TOKEN") or "").strip()
        if not url or not token or len(token) > 512:
            return
        try:
            from tools._macos_bridge import validate_bridge_base_url

            ok, _ = validate_bridge_base_url(url)
        except Exception:
            return
        if not ok:
            return
        if not (env.get("MUGI_AUTH_TOKEN") or "").strip():
            env["MUGI_AUTH_TOKEN"] = token
        if not (env.get("MUGI_CALENDAR_URL") or "").strip():
            env["MUGI_CALENDAR_URL"] = url

    @staticmethod
    def _preexec_limits() -> None:
        try:
            resource.setrlimit(resource.RLIMIT_NOFILE, (256, 256))
        except Exception:
            pass
        try:
            # Best-effort CPU limit for runaway loops.
            resource.setrlimit(resource.RLIMIT_CPU, (120, 120))
        except Exception:
            pass

    def start(self, timeout_s: float = 10.0) -> dict[str, Any]:
        with self._lifecycle_lock:
            if self.state == RuntimeState.QUARANTINED:
                raise RuntimeError(f"runtime is quarantined: {self.quarantined_reason}")
            if self._consecutive_failures > 0:
                time.sleep(min(5.0, self._backoff_s))
            self._stop_locked()
            self._set_state(RuntimeState.STARTING)
            # Child process restarts its frame seq at 1; clear host high-water mark so replay guard accepts it.
            self._last_inbound_seq = 0
            cmd = [
                sys.executable,
                "-m",
                "plugin_runtime_process",
                "--plugin-dir",
                self.plugin_dir,
                "--plugin-slug",
                self.plugin_slug,
                "--plugin-name",
                self.plugin_name,
                "--plugin-hash",
                self.plugin_hash,
                "--auth",
                self._auth,
                "--session-nonce",
                self._session_nonce,
            ]
            self._proc = subprocess.Popen(
                cmd,
                cwd=self.plugin_dir,
                env=self._build_env(),
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=False,
                close_fds=True,
                preexec_fn=self._preexec_limits if os.name != "nt" else None,
            )
            self._stop_evt.clear()
            self._reader_thread = threading.Thread(target=self._reader_loop, daemon=True)
            self._reader_thread.start()
            self._stderr_thread = threading.Thread(target=self._stderr_loop, daemon=True)
            self._stderr_thread.start()
            try:
                _perm = self.permissions or {}
                _raw_sl = _perm.get("shell_allowlist")
                if isinstance(_raw_sl, list):
                    _shell_allowlist = [str(x).strip() for x in _raw_sl if str(x).strip()]
                else:
                    _shell_allowlist = []
                if not _shell_allowlist:
                    _shell_allowlist = list(_DEFAULT_SUBPROCESS_ALLOWLIST)
                self._bootstrap = self.request(
                    ACTION_HANDSHAKE,
                    {
                        "protocol_version": 1,
                        "min_supported": 1,
                        "max_supported": 1,
                        "capabilities": {
                            "network": bool(_perm.get("network", False)),
                            "shell": bool(_perm.get("shell", False)),
                            "shell_allowlist": _shell_allowlist,
                            "filesystem": str(_perm.get("filesystem", "none")),
                            "route_enable": True,
                            "hooks": True,
                            "stream_events": True,
                        },
                        "plugin_hash": self.plugin_hash,
                        "session_nonce": self._session_nonce,
                    },
                    timeout_s=timeout_s,
                )
                if str(self._bootstrap.get("plugin_hash") or "") != str(self.plugin_hash):
                    raise RuntimeError("runtime attestation mismatch: plugin hash")
                if str(self._bootstrap.get("session_nonce") or "") != self._session_nonce:
                    raise RuntimeError("runtime attestation mismatch: session nonce")
                self._set_state(RuntimeState.HEALTHY)
                self._consecutive_failures = 0
                self._backoff_s = 1.0
                return dict(self._bootstrap)
            except Exception:
                self._consecutive_failures += 1
                self._backoff_s = min(8.0, self._backoff_s * 2.0)
                self._set_state(RuntimeState.CRASHED)
                if self._consecutive_failures >= self._max_failures_before_quarantine:
                    self.quarantine("Repeated runtime start failures")
                raise

    def stop(self) -> None:
        with self._lifecycle_lock:
            self._stop_locked()

    def _join_io_threads(self, reader: threading.Thread | None, stderr_thr: threading.Thread | None) -> None:
        timeout = 10.0
        for t in (reader, stderr_thr):
            if t is not None and t.is_alive():
                t.join(timeout=timeout)

    def _stop_locked(self) -> None:
        was_quarantined = self.state == RuntimeState.QUARANTINED
        proc = self._proc
        reader = self._reader_thread
        stderr_thr = self._stderr_thread
        if proc is None:
            if not was_quarantined:
                self._set_state(RuntimeState.STOPPED)
            self._stop_evt.set()
            self._join_io_threads(reader, stderr_thr)
            self._reader_thread = None
            self._stderr_thread = None
            return
        if not was_quarantined:
            self._set_state(RuntimeState.STOPPING)
        try:
            proc.terminate()
            proc.wait(timeout=2.0)
        except Exception:
            try:
                proc.kill()
            except Exception:
                pass
        self.last_exit_code = proc.poll()
        self._proc = None
        if not was_quarantined:
            self._set_state(RuntimeState.STOPPED)
        self._stop_evt.set()
        self._join_io_threads(reader, stderr_thr)
        self._reader_thread = None
        self._stderr_thread = None

    def _stderr_loop(self) -> None:
        proc = self._proc
        if proc is None or proc.stderr is None:
            return
        try:
            while not self._stop_evt.is_set():
                line = proc.stderr.readline()
                if not line:
                    break
                try:
                    text = line.decode("utf-8", errors="replace").strip()
                except Exception:
                    text = str(line).strip()
                if text:
                    self.last_error = text[-2000:]
        except Exception:
            pass

    def _reader_loop(self) -> None:
        proc = self._proc
        if proc is None or proc.stdout is None:
            return
        try:
            while not self._stop_evt.is_set():
                raw = proc.stdout.readline()
                if not raw:
                    break
                frame = decode_frame(raw)
                seq = int(frame.get("seq") or 0)
                if seq > 0 and seq <= self._last_inbound_seq:
                    self.last_error = "inbound replay/non-monotonic IPC seq rejected"
                    continue
                if seq > 0:
                    self._last_inbound_seq = seq
                kind = str(frame.get("kind") or "")
                if kind == "response":
                    rid = str(frame.get("request_id") or "")
                    if rid:
                        with self._pending_lock:
                            q = self._pending.get(rid)
                        if q is not None:
                            q.put(frame)
                elif kind == "event":
                    self._handle_event(frame)
        except Exception as e:
            self.last_error = f"reader loop failed: {e}"
        finally:
            if proc is not None:
                try:
                    exit_code = proc.poll()
                except Exception:
                    exit_code = None
                if exit_code is not None:
                    self.last_exit_code = int(exit_code)
            # Superseded by restart/stop: avoid clobbering state (e.g. CRASHED) after _stop_evt.clear().
            if self._proc is not proc:
                return
            state = self.state
            if self._stop_evt.is_set() or state in {
                RuntimeState.STOPPING,
                RuntimeState.STOPPED,
                RuntimeState.QUARANTINED,
            }:
                return
            prev = self.state
            self._set_state(RuntimeState.CRASHED)
            if self.state == RuntimeState.CRASHED and prev != RuntimeState.CRASHED:
                self.metrics.crash_count += 1

    def _handle_event(self, frame: dict[str, Any]) -> None:
        if frame.get("auth") != self._auth:
            return
        event = str(frame.get("event") or "")
        payload = frame.get("payload") if isinstance(frame.get("payload"), dict) else {}
        if event == "heartbeat":
            self.last_heartbeat_ms = int(time.time() * 1000)
            return
        if event == "settings_set" and self.settings_set:
            key = str(payload.get("key") or "")
            if key:
                self.settings_set(key, payload.get("value"))
            return
        if event in {"tool_stream_progress", "tool_stream_chunk", "tool_stream_final", "cancel_ack"}:
            rid = str(payload.get("request_id") or "")
            if not rid:
                return
            with self._pending_lock:
                cb = self._stream_handlers.get(rid)
            if cb is not None:
                try:
                    cb({"event": event, "payload": payload})
                except Exception:
                    pass
        if event == "stream_event" and self.on_stream_event is not None:
            try:
                self.on_stream_event(payload)
            except Exception:
                pass

    def request(
        self,
        action: str,
        payload: dict[str, Any],
        timeout_s: float = 30.0,
        on_event: Callable[[dict[str, Any]], None] | None = None,
    ) -> dict[str, Any]:
        if self.state == RuntimeState.QUARANTINED:
            raise RuntimeError(self.quarantined_reason or "runtime quarantined")
        proc = self._proc
        if proc is None or proc.stdin is None:
            raise RuntimeError("runtime not started")
        if not self._inflight_sem.acquire(timeout=max(0.1, timeout_s)):
            raise RuntimeError("runtime queue full")
        t0 = time.time()
        req = make_request(
            action=action,
            plugin_id=self.plugin_slug,
            auth=self._auth,
            payload=payload,
            capability_context={"session_nonce": self._session_nonce},
            seq=self._next_seq(),
        )
        rid = req["request_id"]
        q: queue.Queue = queue.Queue(maxsize=1)
        with self._pending_lock:
            self._pending[rid] = q
            if on_event is not None:
                self._stream_handlers[rid] = on_event
        try:
            with self._writer_lock:
                proc.stdin.write(encode_frame(req))
                proc.stdin.flush()
            try:
                resp = q.get(timeout=max(0.1, timeout_s))
            except queue.Empty as e:
                self.metrics.timeout_count += 1
                try:
                    self._send_cancel_request(rid)
                except Exception:
                    pass
                if on_event is not None:
                    try:
                        on_event({"event": EVENT_TOOL_STREAM_TIMEOUT, "payload": {"request_id": rid, "timeout_s": timeout_s}})
                    except Exception:
                        pass
                raise TimeoutError(f"runtime request timeout for action '{action}'") from e
            if not bool(resp.get("ok")):
                err = resp.get("error") if isinstance(resp.get("error"), dict) else {}
                msg = str(err.get("message") or "runtime request failed")
                code = str(err.get("code") or IpcErrorCode.INTERNAL)
                self.last_failure_code = code
                raise RuntimeError(f"{code}: {msg}")
            return resp.get("payload") if isinstance(resp.get("payload"), dict) else {}
        finally:
            with self._pending_lock:
                self._pending.pop(rid, None)
                self._stream_handlers.pop(rid, None)
            self._inflight_sem.release()
            self.metrics.add_invoke((time.time() - t0) * 1000.0)

    def _send_cancel_request(self, request_id: str) -> None:
        proc = self._proc
        if proc is None or proc.stdin is None:
            return
        req = make_request(
            action=ACTION_CANCEL_REQUEST,
            plugin_id=self.plugin_slug,
            auth=self._auth,
            payload={"request_id": str(request_id or "")},
            capability_context={"session_nonce": self._session_nonce},
            seq=self._next_seq(),
        )
        with self._writer_lock:
            proc.stdin.write(encode_frame(req))
            proc.stdin.flush()

    def invoke_tool(
        self,
        *,
        tool_name: str,
        args: dict[str, Any],
        capability_context: dict[str, Any] | None = None,
        timeout_s: float = 60.0,
        on_stream_event: Callable[[dict[str, Any]], None] | None = None,
    ) -> Any:
        payload = {
            "tool_name": str(tool_name or "").strip(),
            "args": args if isinstance(args, dict) else {},
            "capability_context": capability_context if isinstance(capability_context, dict) else {},
        }
        out = self.request(
            "invoke_tool",
            payload,
            timeout_s=timeout_s,
            on_event=on_stream_event,
        )
        return out.get("result")

    def call_hook(
        self,
        *,
        event_name: str,
        args: list[Any],
        kwargs: dict[str, Any] | None = None,
        timeout_s: float = 5.0,
    ) -> None:
        self.request(
            "invoke_hook",
            {
                "event_name": str(event_name or "").strip(),
                "args": list(args or []),
                "kwargs": kwargs if isinstance(kwargs, dict) else {},
            },
            timeout_s=timeout_s,
        )

    def proxy_route(
        self,
        *,
        method: str,
        path: str,
        query: str,
        headers: dict[str, str],
        body: bytes,
        context: dict[str, Any] | None = None,
        timeout_s: float = 90.0,
    ) -> dict[str, Any]:
        return self.request(
            "route_request",
            {
                "method": str(method or "GET"),
                "path": str(path or "/"),
                "query": str(query or ""),
                "headers": {str(k): str(v) for k, v in (headers or {}).items()},
                "body": body.decode("utf-8", errors="replace"),
                "context": context if isinstance(context, dict) else {},
            },
            timeout_s=timeout_s,
        )

    def health_check(self, timeout_s: float = 2.0) -> dict[str, Any]:
        return self.request(ACTION_HEALTH, {}, timeout_s=timeout_s)

    def dispatch_control(self, event: dict[str, Any], timeout_s: float = 30.0) -> dict[str, Any]:
        from plugin_ipc_protocol import ACTION_CONTROL_EVENT
        return self.request(ACTION_CONTROL_EVENT, event, timeout_s=timeout_s)

    def get_setting(self, key: str, default=None):
        if not self.settings_get:
            return default
        return self.settings_get(key, default)

