"""Aggregate B2 microbench cells into a single interference-index sweep summary.

Per cell (variant × prefill_size):
- read metrics.jsonl + run_window.json
- slice the shared engine_*.jsonl by run window
- run interference_index() against the slice
- record (variant, prefill_size, n_overlap, n_clean, tpot_p90_*, idx)
"""

from __future__ import annotations

import argparse
import json
from collections import defaultdict
from pathlib import Path
from typing import Any

from analysis.characterization.joined_analysis import (
    _percentile,
    _vllm_rid_matches,
    interference_index,
    load_engine_state,
    load_jsonl,
    write_json,
)


def _slice_engine_state(
    engine_state_by_worker: dict[str, list[dict]],
    t_start: float,
    t_end: float,
) -> dict[str, list[dict]]:
    sliced: dict[str, list[dict]] = {}
    for worker, steps in engine_state_by_worker.items():
        sliced[worker] = [s for s in steps
                          if t_start <= (s.get("t_unix") or 0.0) <= t_end]
    return sliced


def _to_joined_shape(metrics_rows: list[dict], variant: str) -> list[dict]:
    """Adapt B2 metric rows to what interference_index expects."""
    joined: list[dict] = []
    for r in metrics_rows:
        if r.get("workload") != "decode":
            continue
        joined.append({
            "request_id": r["request_id"],
            "tpot_s": r.get("tpot_s"),
            "ttft_s": r.get("ttft_s"),
            "latency_s": r.get("latency_s"),
            "endpoint_url": r.get("endpoint"),
            "routed_to": r.get("endpoint"),
            "t_first_token_unix": (
                (r["t_dispatch_unix"] + r["ttft_s"])
                if r.get("ttft_s") is not None
                and r.get("t_dispatch_unix") is not None else None
            ),
            "t_finish_unix": r.get("t_finish_unix"),
            "error": r.get("error"),
        })
    return joined


def main() -> None:
    p = argparse.ArgumentParser(description="B2 sweep aggregation")
    p.add_argument("--sweep-dir", type=Path, required=True,
                   help="Top-level dir produced by scripts/b2_interference.py")
    p.add_argument("--engine-state-dir", type=Path, required=True)
    p.add_argument("--worker-map", type=str, required=True,
                   help="URL=worker_id pairs, comma-separated")
    p.add_argument("--out", type=Path, default=None)
    args = p.parse_args()

    worker_map = {}
    for entry in args.worker_map.split(","):
        url, _, wid = entry.strip().partition("=")
        if url and wid:
            worker_map[url] = wid

    engine_state = load_engine_state(args.engine_state_dir)
    rows: list[dict] = []
    for variant_dir in sorted(args.sweep_dir.glob("*/")):
        if variant_dir.name in ("logs",):
            continue
        for cell_dir in sorted(variant_dir.glob("p*/")):
            window_path = cell_dir / "run_window.json"
            metrics_path = cell_dir / "metrics.jsonl"
            if not window_path.exists() or not metrics_path.exists():
                continue
            window = json.loads(window_path.read_text())
            metrics_rows = load_jsonl(metrics_path)
            joined = _to_joined_shape(metrics_rows, variant_dir.name)
            sliced = _slice_engine_state(
                engine_state, window["t_start_unix"], window["t_end_unix"],
            )
            idx = interference_index(joined, sliced, worker_map)
            rows.append({
                "variant": variant_dir.name,
                "prefill_size": int(window["prefill_size"]),
                "decode_endpoint": window["decode_endpoint"],
                "prefill_endpoint": window["prefill_endpoint"],
                "n_decode_requests": sum(1 for r in metrics_rows
                                          if r.get("workload") == "decode"
                                          and r.get("error") is None),
                "n_prefill_injections": sum(1 for r in metrics_rows
                                             if r.get("workload") == "prefill"
                                             and r.get("error") is None),
                **idx,
            })
    summary = {"rows": rows}
    out_path = args.out or args.sweep_dir / "b2_sweep_summary.json"
    write_json(out_path, summary)
    print(json.dumps(rows, indent=2))


if __name__ == "__main__":
    main()