#!/usr/bin/env python3
"""Score fixed-rate request logs and summarize static-vs-oracle frontiers."""

from __future__ import annotations

import argparse
import json
import math
import os
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any, Iterable


TARGET_PASS_RATE = 0.95
TPOT_LIMIT_MS = 50.0
PHASES = ("P01", "P06")
CONFIGS = ("C00", "C10", "C01", "C11")


def atomic_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f"{path.name}.tmp.{os.getpid()}")
    temporary.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n")
    os.replace(temporary, path)


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


def percentile(values: Iterable[float], quantile: float) -> float | None:
    ordered = sorted(float(value) for value in values)
    if not ordered:
        return None
    index = max(0, min(len(ordered) - 1, math.ceil(quantile * len(ordered)) - 1))
    return ordered[index]


def ttft_limit_ms(input_tokens: int) -> float:
    if input_tokens <= 4096:
        return 2000.0
    if input_tokens <= 32768:
        return 4000.0
    return 6000.0


def score_trial(
    request_path: Path,
    result_path: Path,
    *,
    phase: str,
    config: str,
    target_rate: float,
    repetition: int,
    role: str,
) -> dict[str, Any]:
    result = json.loads(result_path.read_text())
    rows = [json.loads(line) for line in request_path.read_text().splitlines() if line]
    clean_start = float(result["warmup_seconds"])
    clean_seconds = float(result["clean_segment_seconds"]) * int(
        result["num_clean_segments"]
    )
    clean_end = clean_start + clean_seconds
    cohort = [row for row in rows if clean_start <= float(row["admitted_s"]) < clean_end]

    ttft_values: list[float] = []
    tpot_values: list[float] = []
    lag_values: list[float] = []
    reasons: Counter[str] = Counter()
    passes = 0
    exact_outputs = 0
    for row in cohort:
        lag_ms = (float(row["admitted_s"]) - float(row["scheduled_s"])) * 1000.0
        lag_values.append(lag_ms)
        request_reasons: list[str] = []
        if not bool(row["success"]):
            request_reasons.append(str(row.get("error_kind") or "request_failed"))
        first = row.get("first_token_s")
        if first is None:
            request_reasons.append("ttft_missing")
            ttft = None
        else:
            ttft = (float(first) - float(row["admitted_s"])) * 1000.0
            ttft_values.append(ttft)
            if ttft > ttft_limit_ms(int(row["input_tokens"])):
                request_reasons.append("ttft_slo")
        actual = row.get("actual_output_tokens")
        requested = int(row["requested_output_tokens"])
        if actual == requested:
            exact_outputs += 1
        if first is None or actual is None or int(actual) <= 1:
            request_reasons.append("tpot_missing")
            tpot = None
        else:
            tpot = (
                (float(row["completed_s"]) - float(first))
                * 1000.0
                / (int(actual) - 1)
            )
            tpot_values.append(tpot)
            if tpot > TPOT_LIMIT_MS:
                request_reasons.append("tpot_slo")
        if request_reasons:
            reasons.update(set(request_reasons))
        else:
            passes += 1

    achieved_rate = len(cohort) / clean_seconds if clean_seconds else 0.0
    pass_rate = passes / len(cohort) if cohort else 0.0
    max_lag_ms = max(lag_values, default=math.inf)
    offered_rate_valid = (
        target_rate > 0 and abs(achieved_rate / target_rate - 1.0) <= 0.05
    )
    schedule_valid = bool(cohort) and max_lag_ms <= 1000.0
    raw_feasible = pass_rate >= TARGET_PASS_RATE
    feasible = raw_feasible and offered_rate_valid and schedule_valid
    invariants = {
        "cohort_nonempty": bool(cohort),
        "clean_duration_positive": clean_seconds > 0,
        "timestamps_nondecreasing": all(
            float(row["scheduled_s"]) <= float(row["admitted_s"])
            <= float(row["completed_s"])
            for row in cohort
        ),
        "exact_output_or_failed": all(
            (not bool(row["success"]))
            or row.get("actual_output_tokens") == row.get("requested_output_tokens")
            for row in cohort
        ),
        "latencies_nonnegative": all(value >= 0 for value in ttft_values + tpot_values),
        "pass_rate_in_0_1": 0.0 <= pass_rate <= 1.0,
        "goodput_nonnegative": passes >= 0,
    }
    if not all(invariants.values()):
        raise RuntimeError(f"trial data invariant failed: {invariants}")
    return {
        "schema": 1,
        "phase": phase,
        "config": config,
        "target_rate_rps": target_rate,
        "repetition": repetition,
        "role": role,
        "clean_start_s": clean_start,
        "clean_end_s": clean_end,
        "clean_seconds": clean_seconds,
        "cohort_n": len(cohort),
        "pass_n": passes,
        "pass_rate": pass_rate,
        "slo_goodput_rps": passes / clean_seconds,
        "achieved_offered_rps": achieved_rate,
        "offered_rate_valid": offered_rate_valid,
        "schedule_valid": schedule_valid,
        "raw_slo_feasible": raw_feasible,
        "feasible": feasible,
        "exact_output_n": exact_outputs,
        "failure_reasons": dict(sorted(reasons.items())),
        "ttft_ms": {
            **numeric(ttft_values),
            "p50": percentile(ttft_values, 0.50),
            "p95": percentile(ttft_values, 0.95),
            "p99": percentile(ttft_values, 0.99),
        },
        "tpot_ms": {
            **numeric(tpot_values),
            "p50": percentile(tpot_values, 0.50),
            "p95": percentile(tpot_values, 0.95),
            "p99": percentile(tpot_values, 0.99),
        },
        "schedule_lag_ms": {
            **numeric(lag_values),
            "p95": percentile(lag_values, 0.95),
            "p99": percentile(lag_values, 0.99),
        },
        "invariants": invariants,
    }


def accepted_rate(rows: list[dict[str, Any]]) -> dict[str, Any]:
    verdicts = [bool(row["feasible"]) for row in rows]
    pass_n = sum(int(row["pass_n"]) for row in rows)
    cohort_n = sum(int(row["cohort_n"]) for row in rows)
    return {
        "trials": len(rows),
        "trial_feasible": verdicts,
        "accepted_feasible": sum(verdicts) > len(verdicts) / 2,
        "pooled_pass_n": pass_n,
        "pooled_cohort_n": cohort_n,
        "pooled_pass_rate": pass_n / cohort_n if cohort_n else 0.0,
        "median_goodput_rps": sorted(float(row["slo_goodput_rps"]) for row in rows)[
            len(rows) // 2
        ],
    }


def frontier_for_cell(rows: list[dict[str, Any]]) -> dict[str, Any]:
    by_rate: dict[float, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        by_rate[float(row["target_rate_rps"])].append(row)
    rates = []
    for rate in sorted(by_rate):
        rates.append({"rate_rps": rate, **accepted_rate(by_rate[rate])})
    verdicts = [bool(row["accepted_feasible"]) for row in rates]
    # Once a failure appears, no higher anchor may pass.
    monotone = not any(
        (not verdicts[i]) and any(verdicts[i + 1 :]) for i in range(len(verdicts))
    )
    feasible_rates = [row["rate_rps"] for row in rates if row["accepted_feasible"]]
    infeasible_rates = [row["rate_rps"] for row in rates if not row["accepted_feasible"]]
    lower = max(feasible_rates, default=None)
    upper_candidates = [rate for rate in infeasible_rates if lower is None or rate > lower]
    upper = min(upper_candidates, default=None)
    boundary_repeated = bool(
        lower is not None
        and upper is not None
        and next(row for row in rates if row["rate_rps"] == lower)["trials"] >= 3
        and next(row for row in rates if row["rate_rps"] == upper)["trials"] >= 3
    )
    bracketed = lower is not None and upper is not None and monotone and boundary_repeated
    return {
        "rates": rates,
        "lower_feasible_rps": lower,
        "upper_infeasible_rps": upper,
        "bracketed": bracketed,
        "boundary_repeated": boundary_repeated,
        "monotone": monotone,
    }


def gap_at_weight(
    lower: dict[str, dict[str, float]],
    upper: dict[str, dict[str, float]],
    p01_weight: float,
) -> dict[str, Any]:
    weights = {"P01": p01_weight, "P06": 1.0 - p01_weight}
    oracle = sum(weights[p] * max(upper[p].values()) for p in PHASES)
    static_values = {
        config: sum(weights[p] * lower[p][config] for p in PHASES)
        for config in CONFIGS
    }
    best_config = max(static_values, key=static_values.get)
    static = static_values[best_config]
    return {
        "p01_weight": p01_weight,
        "oracle_upper_rps": oracle,
        "static_lower_rps": static,
        "best_static_config": best_config,
        "gap": oracle / static - 1.0,
    }


def summarize_trials(rows: list[dict[str, Any]]) -> dict[str, Any]:
    grouped: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        grouped[(str(row["phase"]), str(row["config"]))].append(row)
    expected = {(phase, config) for phase in PHASES for config in CONFIGS}
    if set(grouped) != expected:
        raise RuntimeError(f"cell coverage mismatch: {sorted(set(grouped) ^ expected)}")
    frontiers = {
        phase: {
            config: frontier_for_cell(grouped[(phase, config)])
            for config in CONFIGS
        }
        for phase in PHASES
    }
    all_bracketed = all(
        frontiers[p][c]["bracketed"] for p in PHASES for c in CONFIGS
    )
    all_monotone = all(
        frontiers[p][c]["monotone"] for p in PHASES for c in CONFIGS
    )
    lower = {
        p: {c: float(frontiers[p][c]["lower_feasible_rps"]) for c in CONFIGS}
        for p in PHASES
    } if all_bracketed else {}
    upper = {
        p: {c: float(frontiers[p][c]["upper_infeasible_rps"]) for c in CONFIGS}
        for p in PHASES
    } if all_bracketed else {}
    scan = [gap_at_weight(lower, upper, step / 10000) for step in range(10001)] if all_bracketed else []
    worst = max(scan, key=lambda row: row["gap"]) if scan else None
    equal = gap_at_weight(lower, upper, 0.5) if all_bracketed else None
    distinct_by_cell = {
        f"{p}-{c}": len(
            {round(float(row["slo_goodput_rps"]), 12) for row in grouped[(p, c)]}
        )
        for p in PHASES for c in CONFIGS
    }
    sanity = {
        "trial_count": numeric([row["cohort_n"] for row in rows]),
        "target_rates": numeric([row["target_rate_rps"] for row in rows]),
        "pass_rates": numeric([row["pass_rate"] for row in rows]),
        "goodput_rps": numeric([row["slo_goodput_rps"] for row in rows]),
        "distinct_goodput_by_cell": distinct_by_cell,
        "invariants": {
            "all_counters_nonnegative": all(
                int(row["cohort_n"]) >= 0 and int(row["pass_n"]) >= 0 for row in rows
            ),
            "all_ratios_in_0_1": all(0 <= float(row["pass_rate"]) <= 1 for row in rows),
            "all_trial_invariants": all(all(row["invariants"].values()) for row in rows),
            "all_cells_bracketed": all_bracketed,
            "all_frontiers_monotone": all_monotone,
            "per_cell_results_not_all_identical": all(value > 1 for value in distinct_by_cell.values()),
            "weight_scan_continuous": len(scan) in (0, 10001),
        },
    }
    verdict = "INCONCLUSIVE"
    if all(sanity["invariants"].values()) and worst is not None:
        verdict = "REFUTED" if float(worst["gap"]) < 0.10 else "NOT_ESTABLISHED"
    return {
        "schema": 1,
        "verdict": verdict,
        "threshold": 0.10,
        "frontiers": frontiers,
        "equal_time_conservative": equal,
        "worst_mixture_conservative": worst,
        "sanity": sanity,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    sub = parser.add_subparsers(dest="command", required=True)
    score = sub.add_parser("score")
    score.add_argument("--requests", required=True)
    score.add_argument("--result", required=True)
    score.add_argument("--phase", choices=PHASES, required=True)
    score.add_argument("--config", choices=CONFIGS, required=True)
    score.add_argument("--target-rate", type=float, required=True)
    score.add_argument("--repetition", type=int, required=True)
    score.add_argument("--role", required=True)
    score.add_argument("--out", required=True)
    summary = sub.add_parser("summarize")
    summary.add_argument("--trial-glob", required=True)
    summary.add_argument("--out", required=True)
    args = parser.parse_args()
    if args.command == "score":
        value = score_trial(
            Path(args.requests), Path(args.result), phase=args.phase,
            config=args.config, target_rate=args.target_rate,
            repetition=args.repetition, role=args.role,
        )
    else:
        import glob

        paths = [Path(path) for path in sorted(glob.glob(args.trial_glob, recursive=True))]
        value = summarize_trials([json.loads(path.read_text()) for path in paths])
    atomic_json(Path(args.out), value)
    print(json.dumps(value, sort_keys=True))


if __name__ == "__main__":
    main()