#!/usr/bin/env python3
"""Summarize a one-probe-per-K fresh-engine static-K grid conservatively.

Unlike the frontier helper, this reader treats ``probe_details.jsonl`` as the
source of completion and metric-integrity checks.  A missing or partial probe
is reported as invalid rather than silently summarized.
"""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path
from typing import Any


def load_json(path: Path) -> Any:
    return json.loads(path.read_text(encoding="utf-8"))


def first(path: Path, pattern: str) -> Path | None:
    values = sorted(path.glob(pattern))
    return values[0] if values else None


def num(value: Any) -> float | None:
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        return None
    value = float(value)
    return value if math.isfinite(value) else None


def nonnegative_int(value: Any) -> int | None:
    if isinstance(value, bool) or not isinstance(value, int) or value < 0:
        return None
    return value


def read_one(
    root: Path,
    k: int,
    expected_tokens: int | None,
    expected_request_count: int | None,
) -> dict[str, Any]:
    record: dict[str, Any] = {"k": k, "label": "NoSpec" if k == 0 else f"K={k}"}
    store = root / "stores" / f"k{k}"
    result_path = first(store, "*/trials/trial-*/result.json")
    history_path = first(store, "*/trials/trial-*/probe_history.json")
    details_path = first(store, "*/trials/trial-*/probe_details.jsonl")
    record.update(
        {
            "result_path": str(result_path) if result_path else None,
            "probe_history_path": str(history_path) if history_path else None,
            "probe_details_path": str(details_path) if details_path else None,
        }
    )
    failures: list[str] = []
    if result_path is None or history_path is None or details_path is None:
        failures.append("missing_artifact")
        record.update({"valid": False, "integrity_failures": failures})
        return record

    result = load_json(result_path)
    history = load_json(history_path)
    detail_rows = [json.loads(line) for line in details_path.read_text(encoding="utf-8").splitlines() if line]
    if result.get("status") != "completed":
        failures.append(f"result_status={result.get('status')}")
    if result.get("best_source") != "primary_search":
        failures.append(f"best_source={result.get('best_source')}")
    if result.get("completed_with_probe_failure"):
        failures.append("completed_with_probe_failure")
    if not isinstance(history, list) or len(history) != 1:
        failures.append("history_not_exactly_one_probe")
    if len(detail_rows) != 1:
        failures.append("details_not_exactly_one_probe")
    if failures:
        record.update({"valid": False, "integrity_failures": failures})
        return record

    probe = history[0]
    details = detail_rows[0]
    outcomes = details.get("outcomes") if isinstance(details.get("outcomes"), list) else []
    request_count = nonnegative_int(probe.get("request_count"))
    replayed_request_count = nonnegative_int(probe.get("replayed_request_count"))
    if request_count is None:
        failures.append("invalid_request_count")
        request_count = -1
    if replayed_request_count is None:
        failures.append("invalid_replayed_request_count")
        replayed_request_count = -1
    if details.get("early_stopped") or probe.get("early_stopped"):
        failures.append("early_stopped")
    if len(outcomes) != request_count:
        failures.append(f"outcome_count={len(outcomes)}_expected={request_count}")
    if replayed_request_count != request_count:
        failures.append(
            f"replayed_request_count={replayed_request_count}_expected={request_count}"
        )
    if expected_request_count is not None and request_count != expected_request_count:
        failures.append(
            f"materialized_request_count={expected_request_count}_replayed={request_count}"
        )
    successful = sum(bool(item.get("success")) for item in outcomes if isinstance(item, dict))
    tpot_present = sum(item.get("tpot_ms") is not None for item in outcomes if isinstance(item, dict))
    ttft_present = sum(item.get("ttft_ms") is not None for item in outcomes if isinstance(item, dict))
    nonpositive_tpot = sum(
        isinstance(item, dict)
        and num(item.get("tpot_ms")) is not None
        and num(item.get("tpot_ms")) <= 0
        and (nonnegative_int(item.get("completion_tokens")) or 0) > 1
        for item in outcomes
    )
    tokens_verified = all(
        isinstance(item, dict)
        and item.get("completion_tokens_source") == "usage"
        and item.get("completion_tokens") == expected_tokens
        and item.get("expected_completion_tokens") == expected_tokens
        for item in outcomes
    )
    if successful != request_count:
        failures.append(f"success_count={successful}_expected={request_count}")
    if tpot_present != request_count:
        failures.append(f"tpot_count={tpot_present}_expected={request_count}")
    if ttft_present != request_count:
        failures.append(f"ttft_count={ttft_present}_expected={request_count}")
    if nonpositive_tpot:
        failures.append(f"nonpositive_tpot_for_multi_token_completion={nonpositive_tpot}")
    if expected_tokens is not None and not tokens_verified:
        failures.append("completion_tokens_not_all_usage_verified")
    latency = probe.get("latency_summary") if isinstance(probe.get("latency_summary"), dict) else {}
    tpot = latency.get("tpot_ms") if isinstance(latency.get("tpot_ms"), dict) else {}
    record.update(
        {
            "valid": not failures,
            "integrity_failures": failures,
            "threshold": num(probe.get("threshold")),
            "request_count": request_count,
            "replayed_request_count": replayed_request_count,
            "materialized_request_count": expected_request_count,
            "request_rate": num(probe.get("request_rate")),
            "feasible": bool(probe.get("feasible")),
            "pass_rate": num(probe.get("pass_rate")),
            "success_count": successful,
            "ttft_count": ttft_present,
            "tpot_count": tpot_present,
            "nonpositive_tpot_count": nonpositive_tpot,
            "tpot_ms": {name: num(tpot.get(name)) for name in ("mean", "p50", "p90", "p95", "p99")},
        }
    )
    return record


def data_sanity(records: list[dict[str, Any]]) -> dict[str, Any]:
    valid = [record for record in records if record.get("valid")]
    p95 = [num(record.get("tpot_ms", {}).get("p95")) for record in valid]
    p95 = [value for value in p95 if value is not None]
    pass_rates = [num(record.get("pass_rate")) for record in valid]
    pass_rates = [value for value in pass_rates if value is not None]
    request_rates = [num(record.get("request_rate")) for record in valid]
    request_rates = [value for value in request_rates if value is not None]
    return {
        "n_configurations": len(records),
        "n_valid": len(valid),
        "all_valid": len(valid) == len(records),
        "p95_tpot_ms": {
            "n": len(p95),
            "min": min(p95) if p95 else None,
            "max": max(p95) if p95 else None,
            "distinct_value_count": len(set(p95)),
            "non_negative": all(value >= 0 for value in p95),
        },
        "pass_rate": {
            "n": len(pass_rates),
            "min": min(pass_rates) if pass_rates else None,
            "max": max(pass_rates) if pass_rates else None,
            "distinct_value_count": len(set(pass_rates)),
            "within_0_1": all(0 <= value <= 1 for value in pass_rates),
        },
        "request_rate": {
            "n": len(request_rates),
            "min": min(request_rates) if request_rates else None,
            "max": max(request_rates) if request_rates else None,
            "distinct_value_count": len(set(request_rates)),
            "non_negative": all(value >= 0 for value in request_rates),
        },
        "invariants": {
            "one_record_per_k": len({record["k"] for record in records}) == len(records),
            "all_full_completions": all(
                record.get("success_count") == record.get("request_count")
                for record in records
            ),
            "all_latency_values_present": all(
                record.get("ttft_count") == record.get("request_count")
                and record.get("tpot_count") == record.get("request_count")
                for record in records
            ),
            "all_materialized_request_counts_match": all(
                record.get("materialized_request_count") is None
                or record.get("request_count") == record.get("materialized_request_count")
                for record in records
            ),
            "no_nonpositive_tpot_for_multi_token_completion": all(
                record.get("nonpositive_tpot_count", 0) == 0 for record in records
            ),
        },
    }


def materialized_request_count(manifest: dict[str, Any]) -> int | None:
    windows_path = manifest.get("trace_windows_path")
    window_id = manifest.get("trace_window_id")
    if not isinstance(windows_path, str) or not isinstance(window_id, str):
        return None
    path = Path(windows_path)
    if not path.is_file():
        return None
    payload = load_json(path)
    windows = payload.get("windows") if isinstance(payload, dict) else payload
    if not isinstance(windows, list):
        return None
    for window in windows:
        if not isinstance(window, dict) or window.get("window_id") != window_id:
            continue
        value = window.get("num_requests")
        if isinstance(value, int) and not isinstance(value, bool) and value >= 0:
            return value
    return None


def markdown(records: list[dict[str, Any]], sanity: dict[str, Any]) -> str:
    lines = [
        "# CollectiveSpec fresh-engine fixed-grid summary",
        "",
        "| configuration | valid | feasible | offered req/s | materialized/replayed | completed | pass rate | p95 TPOT (ms) | p99 TPOT (ms) | integrity failures |",
        "|---|---:|---:|---:|---:|---:|---:|---:|---:|---|",
    ]
    for record in records:
        tpot = record.get("tpot_ms") or {}
        fmt = lambda value: "—" if value is None else f"{value:.6g}"
        lines.append(
            "| {label} | {valid} | {feasible} | {rate} | {materialized}/{replayed} | {completed}/{count} | {pass_rate} | {p95} | {p99} | {failures} |".format(
                label=record["label"],
                valid=record.get("valid", False),
                feasible=record.get("feasible", "—"),
                rate=fmt(record.get("request_rate")),
                materialized=record.get("materialized_request_count", "—"),
                replayed=record.get("replayed_request_count", "—"),
                completed=record.get("success_count", "—"),
                count=record.get("request_count", "—"),
                pass_rate=fmt(record.get("pass_rate")),
                p95=fmt(tpot.get("p95")),
                p99=fmt(tpot.get("p99")),
                failures=", ".join(record.get("integrity_failures") or []) or "—",
            )
        )
    lines.extend(["", "## Data sanity", "", "```json", json.dumps(sanity, indent=2, sort_keys=True), "``", ""])
    return "\n".join(lines)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--root", type=Path, required=True)
    parser.add_argument("--output-json", type=Path, required=True)
    parser.add_argument("--output-md", type=Path, required=True)
    args = parser.parse_args()
    manifest = load_json(args.root / "manifest.json")
    expected_tokens = manifest.get("completion_tokens_override")
    if isinstance(expected_tokens, bool) or not isinstance(expected_tokens, int):
        expected_tokens = None
    expected_request_count = materialized_request_count(manifest)
    records = [
        read_one(args.root, int(k), expected_tokens, expected_request_count)
        for k in sorted(int(k) for k in manifest["k_values"])
    ]
    sanity = data_sanity(records)
    output = {"manifest": manifest, "records": records, "data_sanity": sanity}
    args.output_json.write_text(json.dumps(output, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    args.output_md.write_text(markdown(records, sanity), encoding="utf-8")
    print(json.dumps(output, indent=2, sort_keys=True))
    return 0 if sanity["all_valid"] else 2


if __name__ == "__main__":
    raise SystemExit(main())