"""Deterministic, stub-backed Route Session transport benchmark.

This is deliberately a transport harness, not a model benchmark.  It gives
performance work a repeatable baseline without requiring a GPU, a live relay,
or localhost sockets (which are not available in every CI sandbox).
"""

from __future__ import annotations

import argparse
import json
import time
import urllib.request
import zlib
from collections import defaultdict
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Iterable, Literal

TransportMode = Literal["direct", "relay"]
CacheMode = Literal["cached", "stateless"]


@dataclass(frozen=True)
class BenchmarkScenario:
    """Fixed input and expected output for one reproducible Route Session."""

    prompt: str = "Route Session profiling prompt."
    output_tokens: tuple[str, ...] = (" amber", " birch", " cedar", " dogwood")
    activation_bytes: int = 4096
    compression: bool = True


@dataclass(frozen=True)
class SeamSample:
    """One head-to-tail activation transfer, with all durations in milliseconds."""

    phase: Literal["prefill", "decode"]
    token_index: int | None
    session_id: str
    activation_id: str
    seam: str
    mode: TransportMode
    cache_mode: CacheMode
    model_ms: float
    encode_ms: float
    activation_decode_ms: float
    framing_ms: float
    metadata_ms: float
    copy_allocation_ms: float
    copy_allocation_bytes: int
    compression_ms: float
    decompression_ms: float
    connection_setup_ms: float
    queue_wait_ms: float
    local_http_forwarding_ms: float
    transport_ms: float
    seam_latency_ms: float
    payload_bytes: int
    wire_bytes: int
    compression_ratio: float
    connection_attempted: bool


@dataclass(frozen=True)
class BenchmarkRun:
    """JSON-safe result for one mode/cache-mode scenario."""

    scenario: BenchmarkScenario
    mode: TransportMode
    cache_mode: CacheMode
    output_tokens: tuple[str, ...]
    samples: tuple[SeamSample, ...]
    cleanup: dict[str, int | bool]

    def to_dict(self) -> dict:
        samples = [asdict(sample) for sample in self.samples]
        return {
            "scenario": asdict(self.scenario),
            "mode": self.mode,
            "cache_mode": self.cache_mode,
            "output_tokens": list(self.output_tokens),
            "session_id": self.samples[0].session_id if self.samples else "",
            "cleanup": self.cleanup,
            "connections": {
                "attempts": sum(sample.connection_attempted for sample in self.samples),
            },
            "phases": _summaries_by(self.samples, lambda sample: sample.phase),
            "seams": _summaries_by(self.samples, lambda sample: sample.seam),
            "samples": samples,
        }


def _percentile(values: Iterable[float], percentile: float) -> float:
    ordered = sorted(values)
    if not ordered:
        return 0.0
    index = max(0, (len(ordered) * percentile + 99) // 100 - 1)
    return round(ordered[int(index)], 4)


def _summary(samples: list[SeamSample]) -> dict[str, float | int]:
    total_latency_ms = sum(sample.seam_latency_ms for sample in samples)
    return {
        "count": len(samples),
        "p50_latency_ms": _percentile((sample.seam_latency_ms for sample in samples), 50),
        "p95_latency_ms": _percentile((sample.seam_latency_ms for sample in samples), 95),
        "payload_bytes": sum(sample.payload_bytes for sample in samples),
        "wire_bytes": sum(sample.wire_bytes for sample in samples),
        "compression_ratio": round(
            sum(sample.payload_bytes for sample in samples) / max(1, sum(sample.wire_bytes for sample in samples)), 4
        ),
        "connection_attempts": sum(sample.connection_attempted for sample in samples),
        "p50_queue_wait_ms": _percentile((sample.queue_wait_ms for sample in samples), 50),
        "p95_queue_wait_ms": _percentile((sample.queue_wait_ms for sample in samples), 95),
        "tokens_per_sec": round(
            sum(sample.phase == "decode" for sample in samples) / max(0.001, total_latency_ms / 1000), 4
        ),
        "bytes_per_token": round(
            sum(sample.wire_bytes for sample in samples) / max(1, sum(sample.phase == "decode" for sample in samples)), 4
        ),
        "compression_cpu_ms": round(
            sum(sample.compression_ms + sample.decompression_ms for sample in samples), 4
        ),
        "model_execution_ms": round(sum(sample.model_ms for sample in samples), 4),
        "activation_encoding_ms": round(sum(sample.encode_ms for sample in samples), 4),
        "activation_decoding_ms": round(sum(sample.activation_decode_ms for sample in samples), 4),
        "local_http_forwarding_ms": round(sum(sample.local_http_forwarding_ms for sample in samples), 4),
        "peak_buffered_bytes": max((sample.copy_allocation_bytes for sample in samples), default=0),
    }


def _summaries_by(samples: tuple[SeamSample, ...], key) -> dict[str, dict[str, float | int]]:
    groups: dict[str, list[SeamSample]] = defaultdict(list)
    for sample in samples:
        groups[key(sample)].append(sample)
    return {name: _summary(group) for name, group in groups.items()}


class _StubTransport:
    """A deterministic two-node seam with explicit connection ownership."""

    def __init__(self, mode: TransportMode, cache_mode: CacheMode, scenario: BenchmarkScenario) -> None:
        self.mode = mode
        self.cache_mode = cache_mode
        self.scenario = scenario
        self._open_connections: set[str] = set()
        self.session_id = "benchmark-route-session"
        self._activation_count = 0
        self._closed = False

    def transfer(self, phase: Literal["prefill", "decode"], token_index: int | None) -> SeamSample:
        # Cached Route Sessions own one connection per seam in both direct and
        # relay modes. Stateless calls deliberately remain one-shot baselines.
        persistent = self.cache_mode == "cached"
        request_key = "route-session" if persistent else f"{phase}:{token_index}"
        connection_attempted = request_key not in self._open_connections
        self._open_connections.add(request_key)
        self._activation_count += 1

        payload = _activation(self.scenario.activation_bytes, phase, token_index)
        wire = zlib.compress(payload, level=9) if self.scenario.compression else payload
        payload_bytes, wire_bytes = len(payload), len(wire)
        connection_setup_ms = (0.8 if self.mode == "direct" else 1.4) if connection_attempted else 0.0
        queue_wait_ms = 0.0 if self.mode == "direct" else 0.18 + (0.05 if token_index is not None and token_index % 2 else 0.0)
        model_ms = 1.6 if phase == "prefill" else 0.45
        encode_ms = 0.16 if phase == "prefill" else 0.06
        activation_decode_ms = 0.055 if phase == "prefill" else 0.02
        # Keep framing/metadata/copy costs explicit rather than hiding them in
        # serialization or transport time.  The stub owns one binary frame and
        # one response body per hop; no base64 body is modeled.
        framing_ms = 0.035 if phase == "prefill" else 0.012
        metadata_ms = 0.018 if phase == "prefill" else 0.008
        copy_allocation_ms = 0.025 if self.scenario.compression else 0.012
        copy_allocation_bytes = wire_bytes + payload_bytes
        compression_ms = 0.09 if self.scenario.compression else 0.0
        decompression_ms = 0.07 if self.scenario.compression else 0.0
        # Both routes finish by forwarding the decoded activation to the local
        # tail-node HTTP handler; relay adds its own queue before that hop.
        local_http_forwarding_ms = 0.11 if self.mode == "direct" else 0.16
        transport_ms = (0.32 if self.mode == "direct" else 0.61) + wire_bytes / 100_000
        seam_latency_ms = round(
            model_ms + encode_ms + activation_decode_ms + framing_ms + metadata_ms + copy_allocation_ms
            + compression_ms + decompression_ms + connection_setup_ms + queue_wait_ms + transport_ms
            + local_http_forwarding_ms,
            4,
        )
        return SeamSample(
            phase=phase, token_index=token_index, session_id=self.session_id,
            activation_id=f"benchmark-activation-{self._activation_count}", seam="head->tail", mode=self.mode,
            cache_mode=self.cache_mode, model_ms=model_ms, encode_ms=encode_ms,
            activation_decode_ms=activation_decode_ms,
            framing_ms=framing_ms, metadata_ms=metadata_ms,
            copy_allocation_ms=copy_allocation_ms, copy_allocation_bytes=copy_allocation_bytes,
            compression_ms=compression_ms, decompression_ms=decompression_ms,
            connection_setup_ms=connection_setup_ms, queue_wait_ms=queue_wait_ms,
            local_http_forwarding_ms=local_http_forwarding_ms,
            transport_ms=round(transport_ms, 4), seam_latency_ms=seam_latency_ms,
            payload_bytes=payload_bytes, wire_bytes=wire_bytes,
            compression_ratio=round(payload_bytes / wire_bytes, 4), connection_attempted=connection_attempted,
        )

    def close(self) -> dict[str, int | bool]:
        """Close all deterministic owners and expose a CI-checkable snapshot."""
        self._open_connections.clear()
        self._closed = True
        return {
            "session_closed": True,
            "open_connections": 0,
            "queued_activations": 0,
            "telemetry_aggregates": 0,
        }


def _activation(size: int, phase: str, token_index: int | None) -> bytes:
    """Return a compressible but phase-distinguishable activation body."""
    prefix = f"{phase}:{token_index if token_index is not None else 'prompt'}:".encode()
    return (prefix * ((size // len(prefix)) + 1))[:size]


def run_route_session_benchmark(
    mode: TransportMode,
    cache_mode: CacheMode,
    scenario: BenchmarkScenario = BenchmarkScenario(),
) -> BenchmarkRun:
    """Run one fixed two-node prefill + decode Route Session scenario."""
    transport = _StubTransport(mode, cache_mode, scenario)
    try:
        samples = [transport.transfer("prefill", None)]
        samples.extend(transport.transfer("decode", index) for index in range(len(scenario.output_tokens)))
    finally:
        cleanup = transport.close()
    return BenchmarkRun(scenario, mode, cache_mode, scenario.output_tokens, tuple(samples), cleanup)


def run_benchmark_matrix(scenario: BenchmarkScenario = BenchmarkScenario()) -> dict:
    """Run direct/relay and cached/stateless baselines suitable for CI artifacts."""
    runs = [
        run_route_session_benchmark(mode, cache_mode, scenario).to_dict()
        for mode in ("direct", "relay")
        for cache_mode in ("cached", "stateless")
    ]
    return {"schema_version": 1, "runs": runs}


def assert_benchmark(
    run: BenchmarkRun,
    *,
    expected_tokens: Iterable[str],
    expected_connection_attempts: int,
) -> None:
    """Assertion seam for regression tests and future performance gates."""
    assert tuple(expected_tokens) == run.output_tokens, "stub output tokens changed"
    actual_attempts = sum(sample.connection_attempted for sample in run.samples)
    assert actual_attempts == expected_connection_attempts, (
        f"expected {expected_connection_attempts} connections, got {actual_attempts}"
    )


@dataclass(frozen=True)
class PerformanceThresholds:
    """Stable gate limits.

    A cached decode must retain at least a 20% latency/throughput advantage and
    cannot add more than 20% wire bytes per token.  Those deliberately broad
    ratios tolerate ordinary LAN host variance, yet still catch loss of
    connection reuse or a material transport/data-plane slowdown.  Exact
    correctness, ownership, and cleanup invariants are enforced separately.
    """

    max_cached_p50_latency_ratio: float = 0.80
    min_cached_throughput_ratio: float = 1.20
    max_bytes_per_token_ratio: float = 1.20


def assert_performance_gate(
    report: dict,
    *,
    thresholds: PerformanceThresholds = PerformanceThresholds(),
) -> None:
    """Fail CI on a material transport regression, not ordinary host variation.

    The stub's timing is deterministic, but ratios deliberately allow 20% when
    the report is later compared with a LAN capture.  Connection ownership,
    token identity, Route Session stability, and post-run cleanup are exact
    invariants and must never be relaxed.
    """
    runs = {(run["mode"], run["cache_mode"]): run for run in report["runs"]}
    expected = BenchmarkScenario().output_tokens
    for key, run in runs.items():
        assert tuple(run["output_tokens"]) == expected, f"{key}: output tokens changed"
        samples = run["samples"]
        assert len({sample["session_id"] for sample in samples}) == 1, f"{key}: Route Session changed"
        assert len({sample["activation_id"] for sample in samples}) == len(samples), f"{key}: activation IDs reused"
        assert run["cleanup"] == {
            "session_closed": True, "open_connections": 0,
            "queued_activations": 0, "telemetry_aggregates": 0,
        }, f"{key}: resources leaked"
        expected_connections = 1 if key[1] == "cached" else len(samples)
        assert run["connections"]["attempts"] == expected_connections, f"{key}: connection regression"

    for mode in ("direct", "relay"):
        cached = runs[(mode, "cached")]["phases"]["decode"]
        stateless = runs[(mode, "stateless")]["phases"]["decode"]
        assert cached["p50_latency_ms"] <= stateless["p50_latency_ms"] * thresholds.max_cached_p50_latency_ratio, (
            f"{mode}: cached p50 latency regressed"
        )
        assert cached["tokens_per_sec"] >= stateless["tokens_per_sec"] * thresholds.min_cached_throughput_ratio, (
            f"{mode}: cached throughput regressed"
        )
        assert cached["bytes_per_token"] <= stateless["bytes_per_token"] * thresholds.max_bytes_per_token_ratio, (
            f"{mode}: cached bytes/token regressed"
        )


def run_real_model_lan_benchmark(url: str, *, model: str, timeout: float = 120.0) -> dict:
    """Opt-in client-side LAN capture using the same report schema as CI.

    This intentionally makes exactly one OpenAI-compatible request.  It is a
    live validation aid, not a CI input: remote seam CPU/buffer values are zero
    until nodes expose them in a response, while bytes, latency, output and
    connection ownership are measured at the LAN client boundary.
    """
    scenario = BenchmarkScenario()
    body = json.dumps({
        "model": model,
        "messages": [{"role": "user", "content": scenario.prompt}],
        "max_tokens": len(scenario.output_tokens), "temperature": 0,
    }).encode()
    request = urllib.request.Request(
        f"{url.rstrip('/')}/v1/chat/completions", data=body,
        headers={"Content-Type": "application/json", "X-Meshnet-Session": "lan-benchmark-session"}, method="POST",
    )
    started = time.monotonic()
    with urllib.request.urlopen(request, timeout=timeout) as response:
        response_body = response.read()
        session_id = response.headers.get("X-Meshnet-Session", "lan-benchmark-session")
    elapsed_ms = round((time.monotonic() - started) * 1000, 4)
    payload = json.loads(response_body)
    content = payload["choices"][0]["message"]["content"]
    tokens = tuple(content.split())
    sample = SeamSample(
        phase="decode", token_index=0, session_id=session_id, activation_id="lan-activation-1",
        seam="head->tail", mode="direct", cache_mode="cached", model_ms=0.0, encode_ms=0.0,
        activation_decode_ms=0.0,
        framing_ms=0.0, metadata_ms=0.0, copy_allocation_ms=0.0, copy_allocation_bytes=0,
        compression_ms=0.0, decompression_ms=0.0, connection_setup_ms=elapsed_ms,
        queue_wait_ms=0.0, local_http_forwarding_ms=0.0, transport_ms=elapsed_ms, seam_latency_ms=elapsed_ms,
        payload_bytes=len(body), wire_bytes=len(body) + len(response_body), compression_ratio=1.0,
        connection_attempted=True,
    )
    run = BenchmarkRun(
        scenario, "direct", "cached", tokens, (sample,),
        {"session_closed": True, "open_connections": 0, "queued_activations": 0, "telemetry_aggregates": 0},
    )
    return {"schema_version": 1, "source": "real-model-lan-client", "runs": [run.to_dict()]}


def format_summary(report: dict) -> str:
    """Render the compact, human-readable companion to the JSON artifact."""
    lines = ["Route Session benchmark"]
    for run in report["runs"]:
        decode = run["phases"]["decode"]
        seam = run["seams"]["head->tail"]
        lines.append(
            f"{run['mode']:6} {run['cache_mode']:9} "
            f"decode p50/p95 {decode['p50_latency_ms']:.2f}/{decode['p95_latency_ms']:.2f} ms; "
            f"{decode['tokens_per_sec']:.1f} tok/s; {decode['bytes_per_token']:.0f} B/tok; "
            f"seam {seam['payload_bytes']}/{seam['wire_bytes']} B "
            f"({seam['compression_ratio']:.2f}x); connections {run['connections']['attempts']}; "
            f"model/encode/decode {decode['model_execution_ms']:.2f}/"
            f"{decode['activation_encoding_ms']:.2f}/{decode['activation_decoding_ms']:.2f} ms; "
            f"compression {decode['compression_cpu_ms']:.2f} ms; "
            f"HTTP {decode['local_http_forwarding_ms']:.2f} ms; "
            f"queue p95 {decode['p95_queue_wait_ms']:.2f} ms"
        )
    return "\n".join(lines)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Run the deterministic Route Session benchmark")
    parser.add_argument("--json-out", type=Path, help="write the JSON artifact to this path")
    parser.add_argument("--real-model-lan", metavar="URL", help="opt-in OpenAI-compatible LAN endpoint capture")
    parser.add_argument("--model", help="model name required with --real-model-lan")
    parser.add_argument("--timeout", type=float, default=120.0, help="LAN request timeout in seconds")
    parser.add_argument("--no-gate", action="store_true", help="report deterministic results without enforcing thresholds")
    args = parser.parse_args(argv)
    if args.real_model_lan:
        if not args.model:
            parser.error("--model is required with --real-model-lan")
        report = run_real_model_lan_benchmark(args.real_model_lan, model=args.model, timeout=args.timeout)
    else:
        report = run_benchmark_matrix()
        if not args.no_gate:
            assert_performance_gate(report)
    if args.json_out:
        args.json_out.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(format_summary(report))
    return 0


if __name__ == "__main__":  # pragma: no cover - CLI entry point
    raise SystemExit(main())