#!/usr/bin/env python3
"""Exact C1 anchor replay using the pinned AITuner trace/worker/SLO paths."""

from __future__ import annotations

import argparse
import dataclasses
import hashlib
import json
import math
import os
import sys
import time
from pathlib import Path
from typing import Any


AITUNER_ROOT = Path(os.environ.get("AITUNER_ROOT", Path(__file__).resolve().parents[2]))
sys.path.insert(0, str(AITUNER_ROOT / "src"))
os.environ.setdefault("AITUNER_CODEX_BASE_URL", "http://127.0.0.1:1")

from aituner.slo import evaluate_request, summarize_evaluations  # noqa: E402
from aituner.spec import load_study_spec  # noqa: E402
from aituner.trace import load_trace_requests, select_requests_for_threshold  # noqa: E402
from aituner.worker import _probe_drain_deadline, _replay_requests  # noqa: E402


def atomic_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(payload, sort_keys=True, indent=2) + "\n")
    os.replace(tmp, path)


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def numeric(values: list[float | int | None]) -> dict[str, Any]:
    finite = [float(x) for x in values if x is not None and math.isfinite(float(x))]
    return {
        "n": len(values), "finite_n": len(finite), "missing_n": len(values) - len(finite),
        "min": min(finite) if finite else None, "max": max(finite) if finite else None,
        "distinct_n": len(set(finite)),
    }


def load_selected(study_path: Path, anchor: float):
    study = load_study_spec(study_path)
    window, requests = load_trace_requests(study, study_spec_path=study_path)
    selected = select_requests_for_threshold(requests, threshold=anchor)
    return study, window, requests, selected


def selected_summary(selected, duration_s: float, tp: int) -> dict[str, Any]:
    ids = "\n".join(item.row_id for item in selected)
    arrival = "\n".join(f"{item.arrival_s:.12f}" for item in selected)
    lengths = "\n".join(str(item.prompt_tokens_hint) for item in selected)
    return {
        "count": len(selected),
        "offered_req_s": len(selected) / duration_s,
        "offered_req_s_per_gpu": len(selected) / duration_s / tp,
        "request_id_order_sha256": hashlib.sha256(ids.encode()).hexdigest(),
        "arrival_order_sha256": hashlib.sha256(arrival.encode()).hexdigest(),
        "raw_length_order_sha256": hashlib.sha256(lengths.encode()).hexdigest(),
        "arrival_s": numeric([item.arrival_s for item in selected]),
        "raw_input_tokens": numeric([item.prompt_tokens_hint for item in selected]),
        "long_gt4096": sum(int(item.prompt_tokens_hint or 0) > 4096 for item in selected),
    }


def run_replay(args: argparse.Namespace, *, warmup: bool) -> dict[str, Any]:
    study_path = Path(args.study)
    study, window, _requests, selected = load_selected(study_path, args.anchor)
    if warmup:
        first = selected[:16]
        if not any(int(item.prompt_tokens_hint or 0) > 4096 for item in first):
            long_item = next(item for item in selected if int(item.prompt_tokens_hint or 0) > 4096)
            first = [*selected[:15], long_item]
        first = sorted({item.row_id: item for item in first}.values(), key=lambda item: item.arrival_s)
        if len(first) < 16:
            raise RuntimeError("warmup set has fewer than 16 unique requests")
        start = first[0].arrival_s
        selected = [dataclasses.replace(item, arrival_s=item.arrival_s - start) for item in first]

    duration_s = float(window.window_end - window.window_start)
    interval_start_mono_ns = time.monotonic_ns()
    interval_start_wall_ns = time.time_ns()
    outcomes, early_stopped, early_stop_reason = _replay_requests(
        selected,
        base_url=args.base_url,
        timeout_s=study.engine.request_timeout_s,
        max_concurrency=study.trace.max_concurrency,
        target_pass_rate=(0.0 if warmup else study.slo.target_pass_rate),
        max_lag_s=study.trace.early_stop_max_lag_s,
        max_elapsed_s=(
            120.0 if warmup else _probe_drain_deadline(
                selected, study.slo, ceiling=study.trace.early_stop_max_elapsed_s
            )
        ),
        evaluate_outcome=lambda outcome: evaluate_request(outcome, study.slo),
        drain_inflight_on_early_stop=True,
    )
    interval_end_mono_ns = time.monotonic_ns()
    interval_end_wall_ns = time.time_ns()
    evaluations, slo_summary = summarize_evaluations(outcomes, study.slo)
    by_id = {item.row_id: item for item in selected}
    details = []
    for outcome, evaluation in zip(outcomes, evaluations):
        request = by_id[outcome.request_id]
        details.append({
            "request_id": outcome.request_id,
            "sampling_u": request.sampling_u,
            "arrival_s": request.arrival_s,
            "raw_input_tokens": request.prompt_tokens_hint,
            "success": outcome.success,
            "ttft_ms": outcome.ttft_ms,
            "tpot_ms": outcome.tpot_ms,
            "completion_tokens": outcome.completion_tokens,
            "completion_tokens_source": outcome.completion_tokens_source,
            "slo_pass": evaluation.passed,
            "reasons": evaluation.reasons,
            "error": outcome.error,
        })
    out = Path(args.result_dir)
    out.mkdir(parents=True, exist_ok=True)
    with (out / "requests.jsonl").open("w") as f:
        for item in details:
            f.write(json.dumps(item, sort_keys=True) + "\n")
    summary = selected_summary(selected, duration_s, args.tp)
    exact = sum(
        item.success and item.completion_tokens_source == "usage" and item.completion_tokens == 128
        for item in outcomes
    )
    result = {
        "schema": 1,
        "kind": "warmup" if warmup else "anchor",
        "cell": args.cell,
        "anchor": args.anchor,
        "tp": args.tp,
        "mns": args.mns,
        "study_sha256": sha256_file(study_path),
        "interval": {
            "start_mono_ns": interval_start_mono_ns, "end_mono_ns": interval_end_mono_ns,
            "start_wall_ns": interval_start_wall_ns, "end_wall_ns": interval_end_wall_ns,
            "elapsed_s": (interval_end_mono_ns - interval_start_mono_ns) / 1e9,
        },
        "selection": summary,
        "observed_count": len(outcomes),
        "exact_output_count": exact,
        "slo_pass_count": slo_summary["slo_pass_count"],
        "pass_rate": slo_summary["slo_pass_rate"],
        "feasible": bool(slo_summary["feasible"]),
        "early_stopped": early_stopped,
        "early_stop_reason": early_stop_reason,
        "ttft_ms": numeric([item.ttft_ms for item in outcomes]),
        "tpot_ms": numeric([item.tpot_ms for item in outcomes]),
        "invariants": {
            "selected_nonempty": bool(selected),
            "outcomes_cover_selected": len(outcomes) == len(selected),
            "exact_output_or_failed": all(
                (not item.success) or (
                    item.completion_tokens_source == "usage" and item.completion_tokens == 128
                ) for item in outcomes
            ),
            "raw_lengths_present": all(item.prompt_tokens_hint is not None for item in selected),
            "arrival_nondecreasing": all(
                b.arrival_s >= a.arrival_s for a, b in zip(selected, selected[1:])
            ),
            "warmup_16": (len(outcomes) >= 16 if warmup else True),
            "warmup_exact_16": (exact >= 16 if warmup else True),
            "warmup_long": (
                any(int(item.prompt_tokens_hint or 0) > 4096 for item in selected)
                if warmup else True
            ),
        },
    }
    atomic_json(out / "result.json", result)
    print(json.dumps({k: result[k] for k in ("cell", "anchor", "kind", "pass_rate", "feasible")}))
    if not all(result["invariants"].values()):
        raise RuntimeError(f"client invariants failed: {result['invariants']}")
    return result


def preflight(args: argparse.Namespace) -> None:
    ground = json.loads(Path(args.ground_truth).read_text())
    studies = {1: Path(args.primary_study), 2: Path(args.primary_study), 4: Path(args.tp4_study)}
    loaded = {}
    mismatches = []
    values = []
    for cell in ground["cells"]:
        tp = int(cell["tensor_parallel_size"])
        if tp not in loaded:
            _study, _window, requests, _selected = load_selected(studies[tp], 0.0)
            loaded[tp] = requests
        for historical in cell["probe_history"]:
            selected = select_requests_for_threshold(
                loaded[tp], threshold=float(historical["sampling_u"])
            )
            values.append(len(selected))
            if len(selected) != int(historical["request_count"]):
                mismatches.append({
                    "cell": cell["cell_id"], "anchor": historical["sampling_u"],
                    "expected": historical["request_count"], "actual": len(selected),
                })
    result = {
        "schema": 1, "observations": len(values), "mismatches": mismatches,
        "request_counts": numeric(values),
        "invariants": {"observations_92": len(values) == 92, "counts_match": not mismatches},
    }
    atomic_json(Path(args.out), result)
    print(json.dumps(result, sort_keys=True))
    if not all(result["invariants"].values()):
        raise RuntimeError("preflight count reconstruction failed")


def parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser()
    sub = p.add_subparsers(dest="command", required=True)
    pf = sub.add_parser("preflight")
    pf.add_argument("--ground-truth", required=True)
    pf.add_argument("--primary-study", required=True)
    pf.add_argument("--tp4-study", required=True)
    pf.add_argument("--out", required=True)
    for name in ("warmup", "run-anchor"):
        q = sub.add_parser(name)
        q.add_argument("--study", required=True)
        q.add_argument("--cell", required=True)
        q.add_argument("--anchor", type=float, required=True)
        q.add_argument("--tp", type=int, required=True)
        q.add_argument("--mns", type=int, required=True)
        q.add_argument("--base-url", required=True)
        q.add_argument("--result-dir", required=True)
    return p


def main() -> None:
    args = parser().parse_args()
    if args.command == "preflight":
        preflight(args)
    else:
        run_replay(args, warmup=args.command == "warmup")


if __name__ == "__main__":
    main()