#!/usr/bin/env python3
"""Analyze one paired 10-minute code-trace real/sim canary topology."""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
import re
import statistics
from collections import Counter
from pathlib import Path
from typing import Any


METRICS = ("ttft", "tpot", "e2e")
PROM_COUNTERS = ("vllm:prefix_cache_queries_total", "vllm:prefix_cache_hits_total")
ITERATION_RE = re.compile(
    r"Iteration.*?:\s+"
    r"(?P<context_requests>\d+) context requests, "
    r"(?P<context_tokens>\d+) context tokens, "
    r"(?P<generation_requests>\d+) generation requests, "
    r"(?P<generation_tokens>\d+) generation tokens"
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input-root", type=Path, required=True)
    parser.add_argument("--sim-root", type=Path, required=True)
    parser.add_argument("--real-root", type=Path, action="append", required=True)
    parser.add_argument("--topology", required=True)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def percentile(values: list[float], q: float) -> float:
    ordered = sorted(values)
    position = (len(ordered) - 1) * q
    lower = math.floor(position)
    upper = math.ceil(position)
    if lower == upper:
        return ordered[lower]
    return ordered[lower] * (upper - position) + ordered[upper] * (position - lower)


def distribution(values: list[float]) -> dict[str, float | int]:
    if not values:
        raise ValueError("empty distribution")
    return {
        "count": len(values),
        "mean": statistics.fmean(values),
        "p50": percentile(values, 0.5),
        "p90": percentile(values, 0.9),
        "p95": percentile(values, 0.95),
        "p99": percentile(values, 0.99),
        "max": max(values),
    }


def read_csv(path: Path) -> list[dict[str, str]]:
    with path.open(newline="") as stream:
        return list(csv.DictReader(stream))


def find_one(root: Path, name: str) -> Path:
    matches = list(root.glob(f"**/{name}"))
    if len(matches) != 1:
        raise ValueError(f"expected one {name} below {root}, found {matches}")
    return matches[0]


def prom_counter(path: Path, name: str) -> float:
    values = []
    with path.open() as stream:
        for line in stream:
            if line.startswith(name + "{") or line.startswith(name + " "):
                values.append(float(line.rsplit(maxsplit=1)[1]))
    if not values:
        raise ValueError(f"{path}: missing Prometheus counter {name}")
    return sum(values)


def prefix_cache_delta(root: Path) -> dict[str, float]:
    before = root / "metrics/before.prom"
    after = root / "metrics/after.prom"
    deltas = {
        name: prom_counter(after, name) - prom_counter(before, name)
        for name in PROM_COUNTERS
    }
    queries = deltas[PROM_COUNTERS[0]]
    hits = deltas[PROM_COUNTERS[1]]
    if queries <= 0 or hits < 0 or hits > queries:
        raise ValueError(f"{root}: invalid prefix counter deltas {deltas}")
    return {
        "query_tokens": queries,
        "hit_tokens": hits,
        "hit_ratio": hits / queries,
    }


def real_decode_batch(root: Path) -> dict[str, Any]:
    counts: Counter[int] = Counter()
    mixed_steps = 0
    for path in sorted(root.rglob("server.log")):
        with path.open(errors="replace") as stream:
            for line in stream:
                match = ITERATION_RE.search(line)
                if match is None:
                    continue
                context_requests = int(match.group("context_requests"))
                generation_requests = int(match.group("generation_requests"))
                generation_tokens = int(match.group("generation_tokens"))
                if context_requests:
                    mixed_steps += 1
                    continue
                if generation_requests and generation_tokens == generation_requests:
                    counts[generation_requests] += 1

    if not counts:
        return {
            "steps": 0,
            "mixed_steps_excluded": mixed_steps,
            "max": None,
            "share_gt_1": None,
            "histogram": {},
        }
    steps = sum(counts.values())
    return {
        "steps": steps,
        "mixed_steps_excluded": mixed_steps,
        "max": max(counts),
        "share_gt_1": sum(value for key, value in counts.items() if key > 1) / steps,
        "histogram": {str(key): value for key, value in sorted(counts.items())},
    }


def load_sim(root: Path, trace: list[dict[str, str]], trace_sha: str) -> dict[str, Any]:
    manifest = json.loads((root / "manifest.json").read_text())
    if manifest["trace_sha256"] != trace_sha:
        raise ValueError(f"{root}: sim/input trace SHA mismatch")
    rows = read_csv(find_one(root / "metrics", "request_metrics.csv"))
    if len(rows) != len(trace):
        raise ValueError(f"{root}: sim/input request count mismatch")
    values = {
        "ttft": [float(row["ttft"]) for row in rows],
        "tpot": [float(row["tpot"]) for row in rows if row["tpot"].strip()],
        "e2e": [float(row["request_e2e_time"]) for row in rows],
        "waiting": [float(row["request_waiting_time_total"]) for row in rows],
    }
    completions = [
        float(trace_row["arrived_at"]) + float(metric_row["request_e2e_time"]) / 1000
        for trace_row, metric_row in zip(trace, rows)
    ]
    tail_index = max(range(len(completions)), key=completions.__getitem__)
    last_arrival = max(float(row["arrived_at"]) for row in trace)
    summary = json.loads((root / "summary.json").read_text())
    return {
        "values": values,
        "summary": summary,
        "drain": {
            "last_arrival_s": last_arrival,
            "last_completion_s": completions[tail_index],
            "tail_after_last_arrival_s": completions[tail_index] - last_arrival,
            "tail_driver": {
                "request_index": tail_index,
                "arrival_s": float(trace[tail_index]["arrived_at"]),
                "arrival_before_cutoff_s": last_arrival
                - float(trace[tail_index]["arrived_at"]),
                "input_tokens": int(trace[tail_index]["num_prefill_tokens"]),
                "output_tokens": int(trace[tail_index]["num_decode_tokens"]),
                "waiting_ms": values["waiting"][tail_index],
                "e2e_ms": values["e2e"][tail_index],
            },
        },
    }


def load_real(
    root: Path,
    input_manifest: dict[str, Any],
    trace: list[dict[str, str]],
) -> dict[str, Any]:
    result_path = root / "results/result.json"
    result = json.loads(result_path.read_text())
    if result["contract"]["row_vector_sha256"] != input_manifest["paired_row_vector_sha256"]:
        raise ValueError(f"{root}: real/input row digest mismatch")
    requests = result["requests"]
    if len(requests) != len(trace) or not all(row["success"] for row in requests):
        raise ValueError(f"{root}: incomplete or failed real request vector")
    for index, (request, trace_row) in enumerate(zip(requests, trace)):
        observed = (int(request["input_tokens"]), int(request["requested_output_tokens"]))
        expected = (
            int(trace_row["num_prefill_tokens"]),
            int(trace_row["num_decode_tokens"]),
        )
        if observed != expected:
            raise ValueError(f"{root}: request {index} shape {observed} != {expected}")
    values = {
        metric: [
            float(request[f"{metric}_ms"])
            for request in requests
            if request.get(f"{metric}_ms") is not None
        ]
        for metric in METRICS
    }
    completions = [
        float(request["admitted_s"]) + float(request["e2e_ms"]) / 1000
        for request in requests
    ]
    tail_index = max(range(len(completions)), key=completions.__getitem__)
    last_arrival = max(float(request["scheduled_s"]) for request in requests)
    return {
        "root": str(root),
        "result_sha256": sha256(result_path),
        "values": values,
        "summary": result["summary"],
        "prefix_cache": prefix_cache_delta(root),
        "decode_batch": real_decode_batch(root),
        "drain": {
            "last_arrival_s": last_arrival,
            "last_completion_s": completions[tail_index],
            "tail_after_last_arrival_s": completions[tail_index] - last_arrival,
            "tail_driver": {
                "request_index": tail_index,
                "arrival_s": float(requests[tail_index]["scheduled_s"]),
                "arrival_before_cutoff_s": last_arrival
                - float(requests[tail_index]["scheduled_s"]),
                "input_tokens": int(requests[tail_index]["input_tokens"]),
                "output_tokens": int(requests[tail_index]["requested_output_tokens"]),
                "admission_lag_ms": float(requests[tail_index]["admission_lag_ms"]),
                "e2e_ms": float(requests[tail_index]["e2e_ms"]),
            },
        },
    }


def main() -> None:
    args = parse_args()
    input_manifest = json.loads((args.input_root / "manifest.json").read_text())
    trace_path = args.input_root / "frontier.csv"
    trace = read_csv(trace_path)
    if len(trace) != input_manifest["requests"]:
        raise ValueError("input manifest/trace request count mismatch")
    trace_sha = sha256(trace_path)
    sim = load_sim(args.sim_root, trace, trace_sha)
    reals = [load_real(root, input_manifest, trace) for root in args.real_root]
    pooled = {
        metric: [value for real in reals for value in real["values"][metric]]
        for metric in METRICS
    }
    latency = {}
    for metric in METRICS:
        real_dist = distribution(pooled[metric])
        sim_dist = distribution(sim["values"][metric])
        latency[metric] = {
            "real": real_dist,
            "sim": sim_dist,
            "relative_bias_percent": {
                statistic: 100
                * (float(sim_dist[statistic]) - float(real_dist[statistic]))
                / float(real_dist[statistic])
                for statistic in ("mean", "p50", "p90", "p95", "p99")
            },
            "real_per_trial": [
                distribution(real["values"][metric]) for real in reals
            ],
        }
    payload = {
        "schema": "frontier-code-trace-canary-analysis-v1",
        "topology": args.topology,
        "requests_per_trial": len(trace),
        "trials": len(reals),
        "input": {
            "manifest": str(args.input_root / "manifest.json"),
            "paired_row_vector_sha256": input_manifest["paired_row_vector_sha256"],
            "frontier_csv_sha256": trace_sha,
        },
        "latency_ms": latency,
        "prefix_cache": {
            "real_per_trial": [real["prefix_cache"] for real in reals],
            "real_hit_ratio_mean": statistics.fmean(
                real["prefix_cache"]["hit_ratio"] for real in reals
            ),
            "sim": sim["summary"]["prefix_cache"],
        },
        "drain": {
            "interpretation": (
                "Report the max-completion request explicitly; a response that "
                "arrived well before the cutoff can create a long drain tail "
                "without implying queue accumulation."
            ),
            "real_per_trial": [real["drain"] for real in reals],
            "sim": sim["drain"],
        },
        "sim_decode_batch": sim["summary"]["decode_batch"],
        "real_decode_batch_per_trial": [real["decode_batch"] for real in reals],
        "real_artifacts": [
            {
                "root": real["root"],
                "result_sha256": real["result_sha256"],
            }
            for real in reals
        ],
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    print(json.dumps({"output": str(args.output), "topology": args.topology}))


if __name__ == "__main__":
    main()