"""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 platform
import socket
import subprocess
import sys
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())


def _host_manifest() -> dict[str, Any]:
    """Capture non-secret host facts with the report rather than trusting prose."""
    manifest: dict[str, Any] = {
        "hostname": socket.gethostname(),
        "platform": platform.platform(),
        "python": sys.version.split()[0],
        "cpu_count": os.cpu_count(),
    }
    try:
        import torch

        manifest["torch_version"] = torch.__version__
        manifest["cuda_available"] = bool(torch.cuda.is_available())
        if torch.cuda.is_available():
            manifest["accelerator_name"] = torch.cuda.get_device_name(0)
            manifest["accelerator_runtime"] = getattr(torch.version, "cuda", None) or getattr(
                torch.version, "hip", None
            )
    except ImportError:
        manifest["torch_version"] = None
    return manifest


def _validate_config(config: Mapping[str, Any]) -> None:
    """Reject a comparison that could silently mix models or use home storage."""
    try:
        plan = config["plan"]
        root = Path(config["artifact_storage_root"]).resolve(strict=True)
        recipes = config["recipes"]
    except (KeyError, TypeError, OSError) as exc:
        raise BenchmarkError(
            "benchmark config needs an existing artifact_storage_root, plan, and recipes"
        ) from exc
    if not root.is_absolute() or root == Path("/home") or Path("/home") in root.parents:
        raise BenchmarkError("model artifacts must use configured mounted-drive storage, never /home")
    if not isinstance(recipes, list) or not recipes:
        raise BenchmarkError("benchmark config needs at least one recipe")
    for spec in recipes:
        if spec.get("source_model_id") != plan.get("model_id"):
            raise BenchmarkError("every recipe must declare the plan's exact source_model_id")
        if spec.get("source_model_revision") != plan.get("model_revision"):
            raise BenchmarkError("every recipe must declare the plan's exact source_model_revision")
        digest = spec.get("artifact_sha256", "")
        if not isinstance(digest, str) or len(digest) != 64:
            raise BenchmarkError("every recipe must declare its exact 64-character artifact_sha256")
        artifact = Path(spec.get("artifact_path", "")).resolve(strict=True)
        if artifact != root and root not in artifact.parents:
            raise BenchmarkError("every model artifact must be beneath artifact_storage_root")


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", ""),
        source_model_id=spec.get("source_model_id", ""),
        source_model_revision=spec.get("source_model_revision", ""),
        artifact_sha256=spec.get("artifact_sha256", ""),
        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()
    _validate_config(config)
    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={**_host_manifest(), **dict(config.get("host", {}))},
        evidence_class=config.get("evidence_class", "local-real"),
    )