"""Real runtime drivers for the recipe benchmark.

This module is the only place that imports torch, transformers, or spawns a
llama.cpp server, and :mod:`meshnet_node.recipe_benchmark` imports it lazily.
That keeps the default test suite deterministic, GPU-free and download-free
while the real evidence runs through exactly the same measurement core.

Fairness is the whole point of a baseline, so both drivers are held to the same
rules:

* They are handed a **pre-formatted prompt string**.  Neither applies a chat
  template, because a template applied twice — or differently — by two runtimes
  would show up as a speed and drift difference that has nothing to do with the
  runtime.
* They are given the **same CPU thread budget**, so the comparison measures
  kernels rather than how many cores each runtime felt entitled to take.
* They report the runtime's **own prefill/decode split** where it has one, and
  say so honestly where it does not.
"""

from __future__ import annotations

import json
import os
import socket
import subprocess
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any, Mapping

from .recipe_benchmark import (
    BenchmarkError,
    BenchmarkPlan,
    GenerationSample,
    Lane,
    LoadStats,
    PromptSpec,
    RecipeSpec,
    SamplingPolicy,
    build_report,
    measure_recipe,
)

REAL_INFERENCE_ENV = "MESHNET_ENABLE_REAL_INFERENCE_TESTS"


def real_inference_enabled() -> bool:
    """Real runtimes stay off unless the operator opts in explicitly."""
    return os.environ.get(REAL_INFERENCE_ENV) == "1"


def require_real_inference() -> None:
    if not real_inference_enabled():
        raise BenchmarkError(
            f"real model execution is opt-in: set {REAL_INFERENCE_ENV}=1 to run this benchmark"
        )


def _process_rss(pid: int | None = None) -> int:
    """Resident bytes for a process and its children, or 0 when unobservable."""
    try:
        import psutil
    except ImportError:
        return 0
    try:
        process = psutil.Process(pid) if pid else psutil.Process()
        total = process.memory_info().rss
        for child in process.children(recursive=True):
            try:
                total += child.memory_info().rss
            except psutil.Error:
                continue
        return int(total)
    except Exception:
        return 0


def _directory_bytes(path: Path) -> int:
    if path.is_file():
        return path.stat().st_size
    return sum(entry.stat().st_size for entry in path.rglob("*") if entry.is_file())


class TransformersDriver:
    """The current Transformers/safetensors recipe: the correctness reference.

    Generation is a hand-written prefill-then-decode loop rather than
    ``model.generate`` because the benchmark needs the two phases separated: one
    forward over the prompt gives an exact prefill time and TTFT, and the cached
    single-token steps that follow give an exact decode rate.  ``generate`` would
    hand back one blended number.
    """

    def __init__(
        self,
        model_path: str,
        *,
        device: str = "cpu",
        dtype: str = "bfloat16",
        threads: int = 8,
    ) -> None:
        self.model_path = Path(model_path)
        self.device = device
        self.dtype = dtype
        self.threads = threads
        self._model: Any = None
        self._tokenizer: Any = None
        self._torch: Any = None

    def load(self) -> LoadStats:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        self._torch = torch
        torch.set_num_threads(self.threads)
        torch.manual_seed(0)

        started = time.monotonic()
        self._tokenizer = AutoTokenizer.from_pretrained(
            str(self.model_path), local_files_only=True
        )
        self._model = AutoModelForCausalLM.from_pretrained(
            str(self.model_path),
            dtype=getattr(torch, self.dtype),
            local_files_only=True,
        )
        self._model.to(self.device)
        self._model.eval()
        load_ms = (time.monotonic() - started) * 1000

        return LoadStats(
            artifact_bytes=_directory_bytes(self.model_path),
            load_ms=round(load_ms, 4),
            rss_bytes=_process_rss(),
            vram_bytes=self._vram_bytes(),
            backend_detail=(
                f"torch {torch.__version__}; dtype {self.dtype}; "
                f"device {self.device}; intra-op threads {self.threads}"
            ),
        )

    def _vram_bytes(self) -> int:
        torch = self._torch
        if torch is None or self.device == "cpu":
            return 0
        try:
            if torch.cuda.is_available():
                return int(torch.cuda.max_memory_allocated())
        except Exception:
            return 0
        return 0

    def generate(self, prompt: str, sampling: SamplingPolicy) -> GenerationSample:
        if self._model is None:
            raise BenchmarkError("TransformersDriver.generate called before load()")
        torch = self._torch

        # add_special_tokens=False: the plan owns the prompt format, and the
        # llama.cpp recipe is given the identical string.
        encoded = self._tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
        input_ids = encoded["input_ids"].to(self.device)
        prompt_tokens = int(input_ids.shape[-1])
        eos_ids = {self._tokenizer.eos_token_id} | set(
            getattr(self._model.generation_config, "eos_token_id", None) or []
            if isinstance(getattr(self._model.generation_config, "eos_token_id", None), list)
            else []
        )
        eos_ids.discard(None)

        started = time.monotonic()
        with torch.inference_mode():
            outputs = self._model(input_ids=input_ids, use_cache=True)
            past = outputs.past_key_values
            next_id = self._select(outputs.logits[:, -1, :], sampling)
            ttft_ms = (time.monotonic() - started) * 1000

            token_ids = [int(next_id.item())]
            decode_started = time.monotonic()
            while len(token_ids) < sampling.max_output_tokens and token_ids[-1] not in eos_ids:
                outputs = self._model(
                    input_ids=next_id.view(1, 1), past_key_values=past, use_cache=True
                )
                past = outputs.past_key_values
                next_id = self._select(outputs.logits[:, -1, :], sampling)
                token_ids.append(int(next_id.item()))
            decode_ms = (time.monotonic() - decode_started) * 1000

        total_ms = (time.monotonic() - started) * 1000
        emitted = [token for token in token_ids if token not in eos_ids]
        return GenerationSample(
            text=self._tokenizer.decode(emitted, skip_special_tokens=True),
            prompt_tokens=prompt_tokens,
            # The first token is produced by the prefill forward, so the decode
            # rate must not be credited with it.
            decode_tokens=max(0, len(token_ids) - 1),
            ttft_ms=ttft_ms,
            prefill_ms=ttft_ms,
            decode_ms=decode_ms,
            total_ms=total_ms,
        )

    def _select(self, logits: Any, sampling: SamplingPolicy) -> Any:
        if sampling.temperature > 0:
            raise BenchmarkError(
                "this benchmark is greedy-only: sampling noise is indistinguishable from "
                "quantization drift, which is precisely what the quality lane must isolate"
            )
        return logits.argmax(dim=-1)

    def memory_probe(self) -> tuple[int, int]:
        return _process_rss(), self._vram_bytes()

    def close(self) -> None:
        self._model = None
        self._tokenizer = None
        if self._torch is not None:
            import gc

            gc.collect()


def _free_port() -> int:
    with socket.socket() as probe:
        probe.bind(("127.0.0.1", 0))
        return int(probe.getsockname()[1])


class LlamaCppServerDriver:
    """The whole-model llama.cpp/GGUF recipe, driven through ``llama-server``.

    ``llama-server`` is used rather than an in-process binding because it is the
    shape llama.cpp is actually deployed in and the only one that offers
    continuous batching across parallel slots — which is the runtime property
    this project cares about most.  It also reports its own prefill/decode
    timings per request, so the decode rate is the runtime's own number and not
    an inference drawn from a client-side stopwatch.
    """

    def __init__(
        self,
        binary: str,
        gguf_path: str,
        *,
        device: str = "cpu",
        threads: int = 8,
        n_parallel: int = 4,
        context_per_slot: int = 1024,
        n_gpu_layers: int = 0,
        startup_timeout_s: float = 120.0,
    ) -> None:
        self.binary = Path(binary)
        self.gguf_path = Path(gguf_path)
        self.device = device
        self.threads = threads
        self.n_parallel = n_parallel
        self.context_per_slot = context_per_slot
        self.n_gpu_layers = n_gpu_layers
        self.startup_timeout_s = startup_timeout_s
        self._process: subprocess.Popen | None = None
        self._port = 0
        self._log: Any = None

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

    def load(self) -> LoadStats:
        if not self.binary.exists():
            raise BenchmarkError(f"llama-server binary not found at {self.binary}")
        if not self.gguf_path.exists():
            raise BenchmarkError(f"GGUF artifact not found at {self.gguf_path}")

        self._port = _free_port()
        command = [
            str(self.binary),
            "--model", str(self.gguf_path),
            "--host", "127.0.0.1",
            "--port", str(self._port),
            "--threads", str(self.threads),
            "--parallel", str(self.n_parallel),
            # Every slot must hold a whole request, so the pool is sized for the
            # worst case rather than letting llama.cpp silently truncate context.
            "--ctx-size", str(self.context_per_slot * self.n_parallel),
            "--n-gpu-layers", str(self.n_gpu_layers),
            "--no-webui",
        ]
        started = time.monotonic()
        self._log = subprocess.PIPE
        self._process = subprocess.Popen(
            command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
        )
        self._await_health(started)
        load_ms = (time.monotonic() - started) * 1000

        return LoadStats(
            artifact_bytes=self.gguf_path.stat().st_size,
            load_ms=round(load_ms, 4),
            rss_bytes=_process_rss(self._process.pid),
            vram_bytes=0,
            backend_detail=(
                f"llama-server; threads {self.threads}; parallel slots {self.n_parallel}; "
                f"ctx/slot {self.context_per_slot}; gpu layers {self.n_gpu_layers}"
            ),
        )

    def _await_health(self, started: float) -> None:
        while time.monotonic() - started < self.startup_timeout_s:
            if self._process is not None and self._process.poll() is not None:
                raise BenchmarkError(
                    f"llama-server exited with code {self._process.returncode} during startup"
                )
            try:
                with urllib.request.urlopen(f"{self._url}/health", timeout=2) as response:
                    if response.status == 200:
                        return
            except (urllib.error.URLError, OSError):
                time.sleep(0.25)
        raise BenchmarkError(
            f"llama-server did not become healthy within {self.startup_timeout_s:.0f}s"
        )

    def generate(self, prompt: str, sampling: SamplingPolicy) -> GenerationSample:
        if self._process is None:
            raise BenchmarkError("LlamaCppServerDriver.generate called before load()")
        if sampling.temperature > 0:
            raise BenchmarkError("this benchmark is greedy-only; see TransformersDriver._select")

        body = json.dumps({
            "prompt": prompt,
            "n_predict": sampling.max_output_tokens,
            "temperature": 0.0,
            "top_k": 1,
            "top_p": 1.0,
            "seed": sampling.seed,
            # Prompt cache reuse across repeats would measure the cache, not the
            # prefill, and the safetensors recipe has no equivalent.
            "cache_prompt": False,
            "stream": True,
        }).encode()
        request = urllib.request.Request(
            f"{self._url}/completion", data=body,
            headers={"Content-Type": "application/json"}, method="POST",
        )

        started = time.monotonic()
        ttft_ms = 0.0
        chunks: list[str] = []
        timings: Mapping[str, Any] = {}
        with urllib.request.urlopen(request, timeout=600) as response:
            for raw in response:
                line = raw.decode("utf-8").strip()
                if not line.startswith("data:"):
                    continue
                payload = json.loads(line[len("data:"):].strip())
                content = payload.get("content", "")
                if content and not ttft_ms:
                    ttft_ms = (time.monotonic() - started) * 1000
                chunks.append(content)
                if payload.get("stop"):
                    timings = payload.get("timings") or {}
        total_ms = (time.monotonic() - started) * 1000

        if not timings:
            raise BenchmarkError("llama-server returned no timings; cannot report an honest split")

        prefill_ms = float(timings.get("prompt_ms", 0.0))
        decode_ms = float(timings.get("predicted_ms", 0.0))
        return GenerationSample(
            text="".join(chunks),
            prompt_tokens=int(timings.get("prompt_n", 0)),
            decode_tokens=int(timings.get("predicted_n", 0)),
            ttft_ms=ttft_ms or total_ms,
            prefill_ms=prefill_ms,
            decode_ms=decode_ms,
            total_ms=total_ms,
            # Whatever the wall clock saw but the runtime did not attribute to
            # compute is time this request spent waiting for a slot.
            queue_wait_ms=max(0.0, total_ms - prefill_ms - decode_ms),
        )

    def memory_probe(self) -> tuple[int, int]:
        if self._process is None:
            return 0, 0
        return _process_rss(self._process.pid), 0

    def close(self) -> None:
        if self._process is None:
            return
        self._process.terminate()
        try:
            self._process.wait(timeout=20)
        except subprocess.TimeoutExpired:
            self._process.kill()
            self._process.wait(timeout=10)
        self._process = None


def build_driver(spec: Mapping[str, Any], plan: BenchmarkPlan) -> RecipeDriverBundle:
    """Construct the driver named by a recipe's ``driver`` block."""
    driver_spec = dict(spec["driver"])
    kind = driver_spec.pop("type")
    if kind == "transformers":
        return TransformersDriver(**driver_spec)
    if kind == "llama-cpp-server":
        driver_spec.setdefault("n_parallel", max(plan.concurrency_levels))
        return LlamaCppServerDriver(**driver_spec)
    raise BenchmarkError(f"unknown driver type {kind!r}")


RecipeDriverBundle = Any  # a RecipeDriver; named for readability at the call site


def _plan_from_config(config: Mapping[str, Any]) -> BenchmarkPlan:
    raw = config["plan"]
    return BenchmarkPlan(
        plan_id=raw["plan_id"],
        model_id=raw["model_id"],
        model_revision=raw["model_revision"],
        prompts=tuple(PromptSpec(**prompt) for prompt in raw["prompts"]),
        sampling=SamplingPolicy(**raw.get("sampling", {})),
        concurrency_levels=tuple(raw.get("concurrency_levels", (1, 4))),
        repeats=int(raw.get("repeats", 1)),
        warmup_requests=int(raw.get("warmup_requests", 1)),
    )


def _recipe_from_config(spec: Mapping[str, Any]) -> RecipeSpec:
    return RecipeSpec(
        id=spec["id"],
        runtime=spec["runtime"],
        weight_format=spec["weight_format"],
        weight_quantization=spec["weight_quantization"],
        lane=Lane(spec["lane"]),
        device=spec["device"],
        artifact_path=spec.get("artifact_path", ""),
        is_reference=bool(spec.get("is_reference", False)),
        notes=spec.get("notes", ""),
    )


def run_configured_benchmark(config: Mapping[str, Any]) -> dict:
    """Run every recipe in ``config`` against one shared plan and return the report.

    A recipe whose runtime cannot start is recorded as unavailable with the real
    reason rather than dropped: a report that silently omits the recipe that
    crashed would read as a clean result.
    """
    require_real_inference()
    plan = _plan_from_config(config)

    from .recipe_benchmark import RecipeMeasurement  # local import keeps the seam obvious

    measurements = []
    for spec in config["recipes"]:
        recipe = _recipe_from_config(spec)
        try:
            driver = build_driver(spec, plan)
            measurements.append(measure_recipe(driver, recipe, plan))
        except Exception as exc:
            measurements.append(RecipeMeasurement(
                recipe=recipe,
                load=LoadStats(artifact_bytes=0, load_ms=0.0),
                unavailable_reason=f"{type(exc).__name__}: {exc}",
            ))

    return build_report(
        plan,
        measurements,
        host=dict(config.get("host", {})),
        evidence_class=config.get("evidence_class", "local-real"),
    )