"""Runtime adapter interface. Every runtime runs a model in a dedicated local worker process that
speaks the OpenAI HTTP API on 127.0.0.1:<port>. The manager talks to workers uniformly."""

from __future__ import annotations

import asyncio
import os
import signal
import subprocess
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import httpx
import psutil

from ..hardware import process_footprint_bytes
from ..models.estimator import MemoryEstimate, estimate


@dataclass
class WorkerHandle:
    model_id: str
    runtime: str
    port: int
    process: subprocess.Popen
    log_path: Path
    started_at: float = field(default_factory=time.time)
    context: int = 0
    extra: dict[str, Any] = field(default_factory=dict)

    @property
    def base_url(self) -> str:
        return f"http://127.0.0.1:{self.port}"

    @property
    def pid(self) -> int:
        return self.process.pid

    def alive(self) -> bool:
        return self.process.poll() is None

    def memory_bytes(self) -> int:
        if not self.alive():
            return 0
        return process_footprint_bytes(self.process.pid)


class RuntimeAdapter:
    name = "base"

    def __init__(self, settings, log_dir: Path):
        self.settings = settings
        self.log_dir = log_dir

    # ---- to implement ------------------------------------------------------
    def build_command(self, model: dict, port: int, context: int) -> list[str]:
        raise NotImplementedError

    async def is_ready(self, handle: WorkerHandle, client: httpx.AsyncClient) -> tuple[str, str | None]:
        """Return (status, error) where status in loading|ready|error."""
        raise NotImplementedError

    def available(self) -> bool:
        raise NotImplementedError

    # ---- shared -------------------------------------------------------------
    def estimate_memory(self, model: dict, context: int) -> MemoryEstimate:
        return estimate(model["weights_bytes"], self.name, model.get("kv_bytes_per_token") or 0, context,
                        bool(model.get("vision")))

    def spawn(self, model: dict, port: int, context: int) -> WorkerHandle:
        cmd = self.build_command(model, port, context)
        self.log_dir.mkdir(parents=True, exist_ok=True)
        log_path = self.log_dir / f"worker-{model['id']}.log"
        env = dict(os.environ)
        env.setdefault("PYTHONUNBUFFERED", "1")
        env.setdefault("TOKENIZERS_PARALLELISM", "false")
        env.setdefault("HF_HUB_OFFLINE", "1")
        logf = open(log_path, "ab")
        logf.write(f"\n=== {time.strftime('%Y-%m-%d %H:%M:%S')} spawn: {' '.join(cmd)}\n".encode())
        proc = subprocess.Popen(cmd, stdout=logf, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, env=env,
                                start_new_session=True, cwd=str(Path(model["path"]).parent))
        return WorkerHandle(model_id=model["id"], runtime=self.name, port=port, process=proc, log_path=log_path,
                            context=context)

    async def warmup(self, handle: WorkerHandle, client: httpx.AsyncClient, model: dict) -> dict:
        """Tiny inference to prove the model works. Returns timings."""
        t0 = time.time()
        if model.get("embedding"):
            r = await client.post(f"{handle.base_url}/v1/embeddings", json={"model": model["id"], "input": "warm up"},
                                  timeout=600)
            r.raise_for_status()
            data = r.json()
            return {"ttft_ms": round((time.time() - t0) * 1000, 1), "dims": len(data["data"][0]["embedding"])}
        if model.get("reranker"):
            r = await client.post(f"{handle.base_url}/v1/rerank",
                                  json={"model": model["id"], "query": "warm", "documents": ["warm up"]}, timeout=600)
            r.raise_for_status()
            return {"ttft_ms": round((time.time() - t0) * 1000, 1)}
        r = await client.post(f"{handle.base_url}/v1/chat/completions",
                              json={"model": model["id"], "messages": [{"role": "user", "content": "Say OK."}],
                                    "max_tokens": 4, "temperature": 0.0, "stream": False}, timeout=600)
        r.raise_for_status()
        data = r.json()
        text = (data.get("choices") or [{}])[0].get("message", {}).get("content")
        timings = data.get("timings") or {}
        ttft = timings.get("ttft_ms") or timings.get("prompt_ms") or round((time.time() - t0) * 1000, 1)
        return {"ttft_ms": ttft, "text": text, "total_ms": round((time.time() - t0) * 1000, 1)}

    async def stop(self, handle: WorkerHandle, client: httpx.AsyncClient | None = None, grace: float = 8.0) -> None:
        proc = handle.process
        if proc.poll() is not None:
            return
        # 1. polite shutdown
        if client is not None:
            try:
                await client.post(f"{handle.base_url}/shutdown", timeout=2)
            except Exception:
                pass
        # 2. SIGTERM the whole session
        try:
            os.killpg(proc.pid, signal.SIGTERM)
        except Exception:
            try:
                proc.terminate()
            except Exception:
                pass
        t0 = time.time()
        while proc.poll() is None and time.time() - t0 < grace:
            await asyncio.sleep(0.1)
        if proc.poll() is None:
            try:
                os.killpg(proc.pid, signal.SIGKILL)
            except Exception:
                try:
                    proc.kill()
                except Exception:
                    pass
            t0 = time.time()
            while proc.poll() is None and time.time() - t0 < 5:
                await asyncio.sleep(0.1)
        # kill stray children (llama-server forks none, but be safe)
        try:
            for c in psutil.Process(proc.pid).children(recursive=True):
                c.kill()
        except psutil.Error:
            pass