#!/usr/bin/env python3
"""CPU-only Batch 0/1 characterization analyzer.

The script is intentionally conservative about timestamp claims. Actual
per-session in-flight concurrency is only marked available when each analyzed
request has both a dispatch timestamp and a finish/error timestamp. When a run
only has trace timestamps, or only has trace timestamps plus latency, the script
emits an explicit unavailable/estimate status instead of treating it as proof of
online sequentiality.
"""

from __future__ import annotations

import argparse
import collections
import datetime as dt
import hashlib
import json
import math
import platform
import statistics
import subprocess
import sys
from pathlib import Path
from typing import Any


JsonDict = dict[str, Any]


DISPATCH_FIELDS = [
    "dispatch_timestamp_s",
    "start_timestamp_s",
    "request_start_s",
    "t_dispatch",
    "t_proxy_recv",
]
FIRST_TOKEN_FIELDS = [
    "first_token_timestamp_s",
    "t_first_token",
]
FINISH_FIELDS = [
    "finish_timestamp_s",
    "completion_timestamp_s",
    "end_timestamp_s",
    "t_done",
]
ERROR_TIME_FIELDS = [
    "error_timestamp_s",
    "timeout_timestamp_s",
    "t_error",
    "t_timeout",
]
CACHED_TOKEN_FIELDS = [
    "cached_tokens",
    "cache_hit",
    "prefix_cache_hit_tokens",
]


def main() -> None:
    args = parse_args()
    started_at = dt.datetime.now(dt.timezone.utc)

    config = load_optional_json_object(args.config)
    metrics_summary = load_metrics_summary(args.metrics, args.metrics_summary)
    task_name = args.task_name or infer_task_name(args)
    run_date = args.date or dt.datetime.now().strftime("%Y-%m-%d")
    out_dir = args.output_root / run_date / task_name
    prepare_output_dir(out_dir, args.overwrite)
    (out_dir / "raw").mkdir(parents=True, exist_ok=True)
    (out_dir / "figures").mkdir(parents=True, exist_ok=True)

    trace_records, trace_warnings = load_trace_records(args.trace, args.request_limit)
    metrics_records = load_metrics_records(args.metrics)
    breakdown_records = load_breakdown_records(args.breakdown)

    merge_result = merge_records(trace_records, metrics_records, breakdown_records)
    records = merge_result["records"]
    unmatched_metrics = merge_result["unmatched_metrics"]
    unmatched_breakdown = merge_result["unmatched_breakdown"]

    write_jsonl(out_dir / "raw" / "merged_requests.jsonl", records)
    write_jsonl(out_dir / "raw" / "unmatched_metrics.jsonl", unmatched_metrics)
    write_jsonl(out_dir / "raw" / "unmatched_breakdown.jsonl", unmatched_breakdown)

    batch0 = analyze_batch0(records, metrics_summary, config, args.classification)
    batch1 = analyze_batch1(
        records=records,
        unmatched_breakdown=unmatched_breakdown,
        kv_bytes_per_token=args.kv_bytes_per_token,
        block_size=args.block_size,
        shared_prefix_min_sessions=args.shared_prefix_min_sessions,
        system_prefix_blocks=args.system_prefix_blocks,
    )

    write_json(out_dir / "session_concurrency.json", batch0["session_concurrency"])
    write_json(out_dir / "session_arrival_stats.json", batch0["session_arrival_stats"])
    write_json(out_dir / "turn_interval_stats.json", batch0["turn_interval_stats"])
    write_json(out_dir / "trace_profile.json", batch0["trace_profile"])
    (out_dir / "invalid_runs.md").write_text(
        render_invalid_runs(batch0),
        encoding="utf-8",
    )

    write_json(out_dir / "workload_summary.json", batch1["workload_summary"])
    write_json(out_dir / "kv_footprint_summary.json", batch1["kv_footprint_summary"])
    write_json(out_dir / "reuse_decomposition.json", batch1["reuse_decomposition"])
    write_json(out_dir / "session_skew.json", batch1["session_skew"])
    write_json(out_dir / "append_delta_stats.json", batch1["append_delta_stats"])

    figure_status = generate_figures(out_dir, records, batch0, batch1, args.no_figures)
    finished_at = dt.datetime.now(dt.timezone.utc)

    manifest = build_manifest(
        args=args,
        config=config,
        out_dir=out_dir,
        started_at=started_at,
        finished_at=finished_at,
        batch0=batch0,
    )
    manifest["figure_status"] = figure_status
    manifest["input_status"] = {
        "trace_records": len(trace_records),
        "metrics_records": len(metrics_records),
        "breakdown_records": len(breakdown_records),
        "analyzed_records": len(records),
        "unmatched_metrics": len(unmatched_metrics),
        "unmatched_breakdown": len(unmatched_breakdown),
        "trace_warnings": trace_warnings,
        "merge_warnings": merge_result["warnings"],
    }
    write_json(out_dir / "manifest.json", manifest)

    summary = {
        "classification": batch0["classification"],
        "analyzed_records": len(records),
        "manifest": manifest,
        "batch0": {
            "session_concurrency_status": batch0["session_concurrency"]["status"],
            "session_sequential": batch0["session_concurrency"]["session_sequential"],
            "max_inflight_per_session": batch0["session_concurrency"]["max_inflight_per_session"],
            "attempted_requests": batch0["trace_profile"]["attempted_requests"],
            "completed_requests": batch0["trace_profile"]["completed_requests"],
            "error_requests": batch0["trace_profile"]["error_requests"],
        },
        "batch1": {
            "input_stats": batch1["workload_summary"]["input_tokens"],
            "output_stats": batch1["workload_summary"]["output_tokens"],
            "kv_footprint_status": batch1["kv_footprint_summary"]["status"],
            "reuse_status": batch1["reuse_decomposition"]["status"],
            "append_status": batch1["append_delta_stats"]["status"],
        },
        "outputs": relative_output_list(out_dir),
    }
    write_json(out_dir / "summary.json", summary)
    (out_dir / "summary.md").write_text(
        render_summary_md(summary, batch0, batch1),
        encoding="utf-8",
    )
    (out_dir / "audit.md").write_text(
        render_audit_md(batch0, batch1),
        encoding="utf-8",
    )

    print(f"Wrote characterization outputs to {out_dir}")
    print(f"Run classification: {batch0['classification']['label']}")
    print(f"Session sequentiality: {batch0['session_concurrency']['session_sequential']}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Analyze Batch 0/1 workload characterization artifacts.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--trace", type=Path, default=None,
                        help="Trace JSONL path. Supports Ali trace fields used by replayer.trace.")
    parser.add_argument("--metrics", type=Path, default=None,
                        help="Per-request metrics JSONL path.")
    parser.add_argument("--metrics-summary", type=Path, default=None,
                        help="metrics.summary.json path. If omitted, inferred from --metrics when possible.")
    parser.add_argument("--breakdown", type=Path, default=None,
                        help="Proxy breakdown JSON/JSONL path with timing/cache fields.")
    parser.add_argument("--config", type=Path, default=None,
                        help="Optional run config JSON.")
    parser.add_argument("--output-root", type=Path,
                        default=Path("outputs/characterization"),
                        help="Output root. Final path is <root>/<date>/<task_name>.")
    parser.add_argument("--date", default=None,
                        help="Date component for output path. Defaults to local YYYY-MM-DD.")
    parser.add_argument("--task-name", default=None,
                        help="Task/run name for output path.")
    parser.add_argument("--request-limit", type=int, default=None,
                        help="Limit trace records read for a small dry run.")
    parser.add_argument("--policy", default=None,
                        help="Policy label to record in manifest.")
    parser.add_argument("--launch-command", default=None,
                        help="Original launch command to record in manifest.")
    parser.add_argument("--time-scale", type=float, default=None,
                        help="Replay time scale, if known.")
    parser.add_argument("--session-sampling-method", default="",
                        help="Session sampling method label.")
    parser.add_argument("--hash-inputs", action="store_true",
                        help=(
                            "Compute full SHA256 for input files. Disabled by default because "
                            "canonical raw traces can be hundreds of GiB."
                        ))
    parser.add_argument("--kv-bytes-per-token", type=float, default=0.0,
                        help="Model-specific KV bytes per input token. 0 marks KV footprint unavailable.")
    parser.add_argument("--block-size", type=int, default=512,
                        help="Trace hash block size in tokens.")
    parser.add_argument("--shared-prefix-min-sessions", type=int, default=8,
                        help="Hash block session-count threshold for shared/system-prefix reuse.")
    parser.add_argument("--system-prefix-blocks", type=int, default=4,
                        help="Only block positions below this index can be classified as shared/system-prefix.")
    parser.add_argument("--classification", default="auto",
                        choices=[
                            "auto",
                            "online_realistic",
                            "burst_stress",
                            "synthetic_microbench",
                            "invalid_for_online_claim",
                        ],
                        help="Run classification override.")
    parser.add_argument("--gpu-type", default="",
                        help="GPU type for manifest. Leave empty for CPU-only trace analysis.")
    parser.add_argument("--gpu-count", type=int, default=0,
                        help="GPU count for manifest. Use 0 for CPU-only trace analysis.")
    parser.add_argument("--overwrite", action="store_true",
                        help="Overwrite an existing output directory.")
    parser.add_argument("--no-figures", action="store_true",
                        help="Skip matplotlib figure generation.")
    args = parser.parse_args()
    if args.trace is None and args.metrics is None and args.breakdown is None:
        parser.error("provide at least one of --trace, --metrics, or --breakdown")
    if args.block_size <= 0:
        parser.error("--block-size must be positive")
    return args


def prepare_output_dir(out_dir: Path, overwrite: bool) -> None:
    if out_dir.exists() and any(out_dir.iterdir()) and not overwrite:
        raise SystemExit(
            f"output directory already exists and is not empty: {out_dir}\n"
            "pass --overwrite or choose a different --task-name"
        )
    out_dir.mkdir(parents=True, exist_ok=True)


def infer_task_name(args: argparse.Namespace) -> str:
    for path in [args.trace, args.metrics, args.breakdown]:
        if path is not None:
            return path.stem.replace(".", "_")
    return "characterization"


def load_optional_json_object(path: Path | None) -> JsonDict:
    if path is None or not path.exists():
        return {}
    data = json.loads(path.read_text(encoding="utf-8"))
    return data if isinstance(data, dict) else {}


def load_metrics_summary(metrics: Path | None, explicit: Path | None) -> JsonDict:
    candidates: list[Path] = []
    if explicit is not None:
        candidates.append(explicit)
    if metrics is not None:
        if metrics.name == "metrics.jsonl":
            candidates.append(metrics.with_name("metrics.summary.json"))
        else:
            candidates.append(metrics.with_suffix(".summary.json"))
    for path in candidates:
        if path.exists():
            data = json.loads(path.read_text(encoding="utf-8"))
            return data if isinstance(data, dict) else {}
    return {}


def load_trace_records(path: Path | None, request_limit: int | None) -> tuple[list[JsonDict], list[str]]:
    if path is None:
        return [], []
    warnings: list[str] = []
    records: list[JsonDict] = []
    chat_to_session: dict[int, str] = {}
    with path.open("r", encoding="utf-8") as handle:
        for idx, line in enumerate(handle):
            if request_limit is not None and len(records) >= request_limit:
                break
            if not line.strip():
                continue
            row = json.loads(line)
            chat_id_raw = first_present(row, ["chat_id", "id"])
            parent_raw = first_present(row, ["parent_chat_id", "parent_id"], -1)
            chat_id = to_int(chat_id_raw, idx)
            parent_chat_id = to_int(parent_raw, -1)
            session_id = row.get("session_id")
            if session_id is None:
                if parent_chat_id < 0:
                    session_id = str(chat_id)
                else:
                    session_id = chat_to_session.get(parent_chat_id, str(parent_chat_id))
            chat_to_session[chat_id] = str(session_id)

            turn_id = to_int(first_present(row, ["turn", "turn_id"], idx), idx)
            scheduled_dispatch = to_float(first_present(row, ["timestamp", "trace_timestamp_s"]))
            input_tokens = to_int(first_present(row, ["input_length", "input_tokens", "prompt_tokens"]))
            output_tokens = to_int(first_present(row, ["output_length", "output_tokens", "completion_tokens"]))
            hash_ids = row.get("hash_ids", [])
            if not isinstance(hash_ids, list):
                warnings.append(f"row {idx}: hash_ids is not a list; ignoring")
                hash_ids = []

            request_id = str(
                row.get("request_id")
                or f"{session_id}:{turn_id}:{chat_id}:{idx}"
            )
            records.append({
                "request_id": request_id,
                "session_id": str(session_id),
                "turn_id": turn_id,
                "chat_id": chat_id,
                "parent_chat_id": parent_chat_id,
                "scheduled_dispatch_s": scheduled_dispatch,
                "input_tokens": input_tokens,
                "output_tokens": output_tokens,
                "request_type": str(row.get("type", row.get("request_type", ""))),
                "hash_ids": [to_int(h) for h in hash_ids if to_int(h) is not None],
                "trace_present": True,
                "metrics_present": False,
                "breakdown_present": False,
                "source_index": idx,
            })
    return records, warnings


def load_metrics_records(path: Path | None) -> list[JsonDict]:
    if path is None or not path.exists():
        return []
    records: list[JsonDict] = []
    with path.open("r", encoding="utf-8") as handle:
        for idx, line in enumerate(handle):
            if not line.strip():
                continue
            row = json.loads(line)
            records.append(metrics_row_to_record(row, idx))
    return records


def metrics_row_to_record(row: JsonDict, idx: int) -> JsonDict:
    request_id = str(row.get("request_id") or f"metrics:{idx}")
    session_id = row.get("session_id")
    scheduled_dispatch = to_float(first_present(row, ["trace_timestamp_s", "timestamp"]))
    input_tokens = to_int(first_present(row, ["input_length", "input_tokens", "prompt_tokens"]))
    output_tokens = to_int(first_present(row, ["output_length", "requested_output_tokens", "output_tokens"]))
    actual_dispatch = to_float(first_present(row, DISPATCH_FIELDS))
    ttft = to_float(row.get("ttft_s"))
    latency = to_float(row.get("latency_s"))
    first_token = to_float(first_present(row, FIRST_TOKEN_FIELDS))
    finish = to_float(first_present(row, FINISH_FIELDS))
    error_time = to_float(first_present(row, ERROR_TIME_FIELDS))
    if actual_dispatch is not None and first_token is None and ttft is not None:
        first_token = actual_dispatch + ttft
    if actual_dispatch is not None and finish is None and latency is not None:
        finish = actual_dispatch + latency

    prompt_details = row.get("prompt_tokens_details")
    cached_from_details = None
    if isinstance(prompt_details, dict):
        cached_from_details = to_int(prompt_details.get("cached_tokens"))
    cached_tokens = to_int(first_present(row, CACHED_TOKEN_FIELDS, cached_from_details))

    record: JsonDict = {
        "request_id": request_id,
        "session_id": str(session_id) if session_id is not None else None,
        "turn_id": to_int(row.get("turn_id")),
        "scheduled_dispatch_s": scheduled_dispatch,
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "request_type": str(row.get("request_type", "")),
        "hash_ids": [],
        "trace_present": False,
        "metrics_present": True,
        "breakdown_present": False,
        "source_index": idx,
        "effective_input_length": to_int(row.get("effective_input_length")),
        "cached_tokens": cached_tokens,
        "latency_s": latency,
        "ttft_s": ttft,
        "tpot_s": to_float(row.get("tpot_s")),
        "actual_output_tokens": to_int(row.get("actual_output_tokens")),
        "requested_output_tokens": to_int(row.get("requested_output_tokens")),
        "finish_reason": row.get("finish_reason"),
        "error": row.get("error"),
        "actual_dispatch_s": actual_dispatch,
        "actual_first_token_s": first_token,
        "actual_finish_s": finish,
        "actual_error_s": error_time,
    }
    if scheduled_dispatch is not None and latency is not None:
        record["scheduled_finish_estimate_s"] = scheduled_dispatch + latency
    if scheduled_dispatch is not None and ttft is not None:
        record["scheduled_first_token_estimate_s"] = scheduled_dispatch + ttft
    return record


def load_breakdown_records(path: Path | None) -> list[JsonDict]:
    if path is None or not path.exists():
        return []
    text = path.read_text(encoding="utf-8").strip()
    if not text:
        return []
    raw: Any
    if text[0] in "[{":
        raw = json.loads(text)
        if isinstance(raw, dict):
            raw = raw.get("records", [raw])
    else:
        raw = [json.loads(line) for line in text.splitlines() if line.strip()]
    records: list[JsonDict] = []
    for idx, row in enumerate(raw if isinstance(raw, list) else []):
        if not isinstance(row, dict):
            continue
        request_id = str(row.get("request_id") or f"breakdown:{idx}")
        actual_dispatch = to_float(first_present(row, DISPATCH_FIELDS))
        first_token = to_float(first_present(row, FIRST_TOKEN_FIELDS))
        finish = to_float(first_present(row, FINISH_FIELDS))
        error_time = to_float(first_present(row, ERROR_TIME_FIELDS))
        cached_tokens = to_int(first_present(row, CACHED_TOKEN_FIELDS))
        input_tokens = to_int(first_present(row, ["input_length", "input_tokens", "prompt_tokens"]))
        estimated_new = to_int(first_present(row, ["estimated_new_tokens", "uncached_tokens", "new_tokens"]))
        records.append({
            "request_id": request_id,
            "session_id": str(row["session_id"]) if "session_id" in row else None,
            "turn_id": to_int(first_present(row, ["turn_id", "turn"])),
            "scheduled_dispatch_s": to_float(first_present(row, ["trace_timestamp_s", "timestamp"])),
            "input_tokens": input_tokens,
            "output_tokens": to_int(first_present(row, ["output_length", "output_tokens"])),
            "hash_ids": [],
            "trace_present": False,
            "metrics_present": False,
            "breakdown_present": True,
            "source_index": idx,
            "cached_tokens": cached_tokens,
            "estimated_new_tokens": estimated_new,
            "actual_dispatch_s": actual_dispatch,
            "actual_first_token_s": first_token,
            "actual_finish_s": finish,
            "actual_error_s": error_time,
            "route_class": row.get("route_class"),
            "policy": row.get("policy"),
            "routed_to": row.get("routed_to"),
        })
    return records


def merge_records(
    trace_records: list[JsonDict],
    metrics_records: list[JsonDict],
    breakdown_records: list[JsonDict],
) -> JsonDict:
    warnings: list[str] = []
    canonical = [dict(r) for r in (trace_records or metrics_records or breakdown_records)]
    by_id: dict[str, JsonDict] = {str(r["request_id"]): r for r in canonical}

    unmatched_metrics: list[JsonDict] = []
    if trace_records:
        for metric in metrics_records:
            target = by_id.get(str(metric["request_id"]))
            if target is None:
                unmatched_metrics.append(metric)
                continue
            overlay_record(target, metric, "metrics")
    else:
        unmatched_metrics = []

    unmatched_breakdown: list[JsonDict] = []
    if trace_records or metrics_records:
        for item in breakdown_records:
            target = by_id.get(str(item["request_id"]))
            if target is None:
                unmatched_breakdown.append(item)
                continue
            overlay_record(target, item, "breakdown")
    else:
        unmatched_breakdown = []

    if trace_records and metrics_records and len(unmatched_metrics) == len(metrics_records):
        warnings.append(
            "no metrics request_id matched trace request_id; metrics were saved under raw/unmatched_metrics.jsonl"
        )
    if (trace_records or metrics_records) and breakdown_records and len(unmatched_breakdown) == len(breakdown_records):
        warnings.append(
            "no breakdown request_id matched analyzed records; breakdown rows were saved under raw/unmatched_breakdown.jsonl"
        )

    return {
        "records": canonical,
        "unmatched_metrics": unmatched_metrics,
        "unmatched_breakdown": unmatched_breakdown,
        "warnings": warnings,
    }


def overlay_record(target: JsonDict, source: JsonDict, source_name: str) -> None:
    target[f"{source_name}_present"] = True
    for key, value in source.items():
        if key in {"request_id", "source_index", "trace_present", "metrics_present", "breakdown_present"}:
            continue
        if value is None:
            continue
        if key == "hash_ids" and not value:
            continue
        if target.get(key) is None or key in {
            "cached_tokens",
            "latency_s",
            "ttft_s",
            "tpot_s",
            "actual_output_tokens",
            "requested_output_tokens",
            "finish_reason",
            "error",
            "actual_dispatch_s",
            "actual_first_token_s",
            "actual_finish_s",
            "actual_error_s",
            "scheduled_finish_estimate_s",
            "scheduled_first_token_estimate_s",
            "estimated_new_tokens",
            "route_class",
            "policy",
            "routed_to",
        }:
            target[key] = value


def analyze_batch0(
    records: list[JsonDict],
    metrics_summary: JsonDict,
    config: JsonDict,
    classification_override: str,
) -> JsonDict:
    trace_profile = build_trace_profile(records, metrics_summary)
    actual_concurrency = compute_session_concurrency(
        records,
        start_key="actual_dispatch_s",
        finish_key="actual_finish_s",
        error_key="actual_error_s",
    )
    scheduled_estimate = compute_session_concurrency(
        records,
        start_key="scheduled_dispatch_s",
        finish_key="scheduled_finish_estimate_s",
        error_key=None,
    )

    session_concurrency = {
        **actual_concurrency,
        "timestamp_requirements": {
            "dispatch_timestamp_fields": DISPATCH_FIELDS,
            "finish_timestamp_fields": FINISH_FIELDS,
            "error_timestamp_fields": ERROR_TIME_FIELDS,
            "note": (
                "Actual max in-flight is only conclusive when dispatch and finish/error "
                "timestamps exist per request. scheduled_estimate uses trace_timestamp_s "
                "plus latency_s and is not proof of replay sequentiality."
            ),
        },
        "scheduled_estimate": scheduled_estimate,
    }

    arrival_stats = compute_session_arrivals(records)
    turn_interval_stats = compute_turn_intervals(records)
    classification = classify_run(
        override=classification_override,
        session_concurrency=session_concurrency,
        trace_profile=trace_profile,
        config=config,
    )
    return {
        "session_concurrency": session_concurrency,
        "session_arrival_stats": arrival_stats,
        "turn_interval_stats": turn_interval_stats,
        "trace_profile": trace_profile,
        "classification": classification,
    }


def build_trace_profile(records: list[JsonDict], metrics_summary: JsonDict) -> JsonDict:
    records_with_session = [r for r in records if r.get("session_id") is not None]
    sessions = group_by_session(records)
    scheduled = sorted(clean_numbers(r.get("scheduled_dispatch_s") for r in records))
    input_tokens = clean_numbers(r.get("input_tokens") for r in records)
    output_tokens = clean_numbers(r.get("output_tokens") for r in records)
    request_types = collections.Counter(str(r.get("request_type", "")) for r in records)
    errors = [r for r in records if r.get("error")]
    metrics_present = [r for r in records if r.get("metrics_present")]
    successes = [r for r in metrics_present if not r.get("error")]
    timing_completed = [
        r for r in records
        if to_float(r.get("actual_finish_s")) is not None
    ]
    timing_errors = [
        r for r in records
        if to_float(r.get("actual_error_s")) is not None
    ]
    completed_count = len(successes) if metrics_present else (
        len(timing_completed) if timing_completed or timing_errors else None
    )
    error_count = len(errors) if metrics_present else (
        len(timing_errors) if timing_completed or timing_errors else None
    )
    wall_clock = to_float(metrics_summary.get("wall_clock_s"))
    if wall_clock is None and scheduled:
        wall_clock = scheduled[-1] - scheduled[0]
    if wall_clock is None:
        actual_starts = clean_numbers(r.get("actual_dispatch_s") for r in records)
        actual_ends = clean_numbers(
            first_present(r, ["actual_finish_s", "actual_error_s"]) for r in records
        )
        if actual_starts and actual_ends:
            wall_clock = max(actual_ends) - min(actual_starts)
    latency_stats = stats(
        [to_float(r.get("latency_s")) for r in successes if to_float(r.get("latency_s")) is not None]
    )
    ttft_stats = stats(
        [to_float(r.get("ttft_s")) for r in successes if to_float(r.get("ttft_s")) is not None]
    )
    tpot_stats = stats(
        [to_float(r.get("tpot_s")) for r in successes if to_float(r.get("tpot_s")) is not None]
    )
    hash_counts = [len(r.get("hash_ids") or []) for r in records]
    unique_hashes: set[int] = set()
    hash_refcount: collections.Counter[int] = collections.Counter()
    for r in records:
        for hid in r.get("hash_ids") or []:
            unique_hashes.add(int(hid))
            hash_refcount[int(hid)] += 1
    shared_hashes = sum(1 for count in hash_refcount.values() if count > 1)
    session_turns = [len(v) for v in sessions.values()]
    return {
        "attempted_requests": len(records),
        "completed_requests": completed_count,
        "error_requests": error_count,
        "metrics_available": bool(metrics_present),
        "records_with_session": len(records_with_session),
        "session_count": len(sessions),
        "multi_turn_sessions": sum(1 for turn_count in session_turns if turn_count > 1),
        "turns_per_session": stats(session_turns),
        "trace_span_s": (scheduled[-1] - scheduled[0]) if len(scheduled) >= 2 else 0.0,
        "request_rate_per_s": len(records) / (scheduled[-1] - scheduled[0]) if len(scheduled) >= 2 and scheduled[-1] > scheduled[0] else None,
        "goodput_per_s": completed_count / wall_clock if completed_count is not None and wall_clock and wall_clock > 0 else None,
        "latency_stats_s": latency_stats,
        "ttft_stats_s": ttft_stats,
        "tpot_stats_s": tpot_stats,
        "latency_percentiles_over_successes_only": bool(metrics_present),
        "input_tokens": stats(input_tokens),
        "output_tokens": stats(output_tokens),
        "request_type_counts": dict(request_types),
        "hash_block_stats": {
            "records_with_hash_ids": sum(1 for count in hash_counts if count > 0),
            "unique_hash_blocks": len(unique_hashes),
            "shared_hash_blocks": shared_hashes,
            "shared_hash_block_fraction": shared_hashes / len(unique_hashes) if unique_hashes else None,
            "hash_blocks_per_request": stats(hash_counts),
        },
        "comparison_metric_availability": {
            "per_worker_queue_metrics": {
                "status": "unavailable",
                "needed_input": "per-worker queue depth/delay metrics exported by proxy or engine",
            },
            "per_worker_gpu_utilization": {
                "status": "unavailable",
                "needed_input": "gpu_util.csv or equivalent per-worker utilization artifact",
            },
            "per_worker_kv_occupancy": {
                "status": "unavailable",
                "needed_input": "per-worker KV occupancy metrics from engine /metrics snapshots",
            },
            "per_worker_apc_cache_hit": {
                "status": (
                    "aggregate_available"
                    if "prefix_cache_hit_ratio" in metrics_summary
                    else "unavailable"
                ),
                "needed_input": "per-worker APC/cache-hit counters; metrics.summary.json only provides aggregate APC when present",
            },
        },
        "metrics_summary_fields": sorted(metrics_summary.keys()),
    }


def compute_session_concurrency(
    records: list[JsonDict],
    *,
    start_key: str,
    finish_key: str,
    error_key: str | None,
) -> JsonDict:
    records_with_session = [r for r in records if r.get("session_id") is not None]
    intervals: dict[str, list[tuple[float, float, str]]] = collections.defaultdict(list)
    missing_start = 0
    missing_finish = 0
    for r in records_with_session:
        start = to_float(r.get(start_key))
        finish = to_float(r.get(finish_key))
        if finish is None and error_key is not None:
            finish = to_float(r.get(error_key))
        if start is None:
            missing_start += 1
            continue
        if finish is None:
            missing_finish += 1
            continue
        if finish < start:
            finish = start
        intervals[str(r["session_id"])].append((start, finish, str(r.get("request_id", ""))))

    per_session: list[JsonDict] = []
    max_seen: int | None = None
    violations: list[JsonDict] = []
    overlap_examples: list[JsonDict] = []
    for sid, items in intervals.items():
        events: list[tuple[float, int]] = []
        for start, finish, _ in items:
            events.append((start, 1))
            events.append((finish, -1))
        events.sort(key=lambda item: (item[0], item[1]))
        cur = 0
        max_cur = 0
        for _, delta in events:
            cur += delta
            max_cur = max(max_cur, cur)
        max_seen = max(max_seen or 0, max_cur)
        sorted_items = sorted(items)
        session_overlaps = 0
        prev_finish = None
        prev_id = None
        for start, finish, req_id in sorted_items:
            if prev_finish is not None and start < prev_finish:
                session_overlaps += 1
                if len(overlap_examples) < 20:
                    overlap_examples.append({
                        "session_id": sid,
                        "previous_request_id": prev_id,
                        "request_id": req_id,
                        "previous_finish_s": prev_finish,
                        "request_start_s": start,
                        "overlap_s": prev_finish - start,
                    })
            if prev_finish is None or finish > prev_finish:
                prev_finish = finish
                prev_id = req_id
        item = {
            "session_id": sid,
            "interval_count": len(items),
            "max_inflight": max_cur,
            "overlap_count": session_overlaps,
        }
        per_session.append(item)
        if max_cur > 1:
            violations.append(item)

    covered = sum(len(v) for v in intervals.values())
    if not records_with_session:
        status = "unavailable"
    elif covered == len(records_with_session):
        status = "available"
    elif covered > 0:
        status = "partial"
    else:
        status = "unavailable"

    session_sequential = None
    if status == "available":
        session_sequential = bool((max_seen or 0) <= 1)
    elif status == "partial" and max_seen and max_seen > 1:
        session_sequential = False

    per_session.sort(key=lambda r: (-int(r["max_inflight"]), -int(r["interval_count"]), str(r["session_id"])))
    return {
        "status": status,
        "start_field": start_key,
        "finish_field": finish_key,
        "records_with_session": len(records_with_session),
        "intervals_with_timestamps": covered,
        "missing_dispatch_timestamps": missing_start,
        "missing_finish_or_error_timestamps": missing_finish,
        "session_sequential": session_sequential,
        "max_inflight_per_session": max_seen if status in {"available", "partial"} else None,
        "sessions_with_overlap": len(violations),
        "per_session": per_session,
        "overlap_examples": overlap_examples,
        "unavailable_reason": (
            None if status == "available"
            else "missing dispatch and/or finish/error timestamps for at least one request"
        ),
    }


def compute_session_arrivals(records: list[JsonDict]) -> JsonDict:
    sessions = group_by_session(records)
    starts: list[float] = []
    turns_per_session: list[int] = []
    for rows in sessions.values():
        times = clean_numbers(r.get("scheduled_dispatch_s") for r in rows)
        if not times:
            continue
        starts.append(min(times))
        turns_per_session.append(len(rows))
    starts.sort()
    if not starts:
        return {
            "status": "unavailable",
            "reason": "missing session_id or scheduled_dispatch_s/timestamp",
        }
    base = starts[0]
    offsets = [v - base for v in starts]
    inter_arrivals = [starts[i + 1] - starts[i] for i in range(len(starts) - 1)]
    return {
        "status": "available",
        "session_count": len(starts),
        "session_start_offsets_s": offsets,
        "session_start_offset_stats_s": stats(offsets),
        "session_interarrival_stats_s": stats(inter_arrivals),
        "arrival_span_s": starts[-1] - starts[0] if len(starts) > 1 else 0.0,
        "mean_session_arrival_rate_per_s": len(starts) / (starts[-1] - starts[0]) if len(starts) > 1 and starts[-1] > starts[0] else None,
        "turns_per_session": stats(turns_per_session),
    }


def compute_turn_intervals(records: list[JsonDict]) -> JsonDict:
    sessions = group_by_session(records)
    scheduled_intervals: list[float] = []
    actual_gaps: list[float] = []
    missing_actual = 0
    for rows in sessions.values():
        rows = sorted(
            rows,
            key=lambda r: (
                value_or_inf(r.get("turn_id")),
                value_or_inf(r.get("scheduled_dispatch_s")),
                str(r.get("request_id", "")),
            ),
        )
        for prev, cur in zip(rows, rows[1:]):
            prev_start = to_float(prev.get("scheduled_dispatch_s"))
            cur_start = to_float(cur.get("scheduled_dispatch_s"))
            if prev_start is not None and cur_start is not None:
                scheduled_intervals.append(cur_start - prev_start)
            prev_finish = to_float(prev.get("actual_finish_s"))
            cur_actual_start = to_float(cur.get("actual_dispatch_s"))
            if prev_finish is not None and cur_actual_start is not None:
                actual_gaps.append(cur_actual_start - prev_finish)
            else:
                missing_actual += 1
    return {
        "status": "available" if scheduled_intervals else "unavailable",
        "scheduled_turn_intervals_s": stats(scheduled_intervals),
        "scheduled_turn_interval_values_s": scheduled_intervals,
        "actual_turn_gaps_status": "available" if actual_gaps and missing_actual == 0 else "partial" if actual_gaps else "unavailable",
        "actual_turn_gaps_s": stats(actual_gaps),
        "actual_negative_gap_count": sum(1 for v in actual_gaps if v < 0),
        "missing_actual_gap_pairs": missing_actual,
        "note": "scheduled intervals use trace timestamps; actual gaps require dispatch and finish timestamps.",
    }


def classify_run(
    *,
    override: str,
    session_concurrency: JsonDict,
    trace_profile: JsonDict,
    config: JsonDict,
) -> JsonDict:
    if override != "auto":
        return {
            "label": override,
            "source": "user_override",
            "reason": "classification override was provided",
        }
    time_scale = to_float(config.get("time_scale"))
    max_sessions = first_present(config, ["max_sessions", "max_inflight_sessions"])
    stress_indicators = []
    if time_scale is not None and not math.isclose(time_scale, 1.0):
        stress_indicators.append(f"time_scale={time_scale}")
    if max_sessions is not None:
        stress_indicators.append(f"max_sessions={max_sessions}")

    sequential = session_concurrency.get("session_sequential")
    status = session_concurrency.get("status")
    if sequential is True and not stress_indicators:
        label = "online_realistic"
        reason = "actual per-session interval timestamps prove max_inflight_per_session <= 1"
    elif sequential is False:
        label = "invalid_for_online_claim"
        reason = "per-session interval overlap was observed"
    elif stress_indicators:
        label = "burst_stress"
        reason = "stress indicators present and sequentiality is not proven: " + ", ".join(stress_indicators)
    elif trace_profile.get("attempted_requests", 0) <= 10:
        label = "synthetic_microbench"
        reason = "small run without enough timestamp data to support online workload claims"
    elif status != "available":
        label = "invalid_for_online_claim"
        reason = "actual dispatch/finish timestamps are unavailable, so online sequentiality cannot be proven"
    else:
        label = "invalid_for_online_claim"
        reason = "run does not satisfy online-realistic criteria"
    return {
        "label": label,
        "source": "auto",
        "reason": reason,
        "stress_indicators": stress_indicators,
    }


def analyze_batch1(
    *,
    records: list[JsonDict],
    unmatched_breakdown: list[JsonDict],
    kv_bytes_per_token: float,
    block_size: int,
    shared_prefix_min_sessions: int,
    system_prefix_blocks: int,
) -> JsonDict:
    input_tokens = [v for v in clean_numbers(r.get("input_tokens") for r in records)]
    output_tokens = [v for v in clean_numbers(r.get("output_tokens") for r in records)]
    ratios = [
        inp / out
        for inp, out in zip(
            [to_float(r.get("input_tokens")) for r in records],
            [to_float(r.get("output_tokens")) for r in records],
        )
        if inp is not None and out is not None and out > 0
    ]
    workload_summary = {
        "status": "available" if input_tokens or output_tokens else "unavailable",
        "request_count": len(records),
        "input_tokens": stats(input_tokens),
        "output_tokens": stats(output_tokens),
        "input_output_ratio": stats(ratios),
        "total_input_tokens": sum(input_tokens),
        "total_output_tokens": sum(output_tokens),
        "long_input_thresholds": {
            "requests_gt_8k": count_gt(input_tokens, 8_000),
            "requests_gt_32k": count_gt(input_tokens, 32_000),
            "requests_gt_64k": count_gt(input_tokens, 64_000),
        },
    }

    kv_footprint_summary = build_kv_footprint_summary(input_tokens, kv_bytes_per_token)
    append_delta_stats = build_append_delta_stats(records, unmatched_breakdown)
    session_skew = build_session_skew(records)
    reuse_decomposition = build_reuse_decomposition(
        records=records,
        block_size=block_size,
        shared_prefix_min_sessions=shared_prefix_min_sessions,
        system_prefix_blocks=system_prefix_blocks,
    )
    return {
        "workload_summary": workload_summary,
        "kv_footprint_summary": kv_footprint_summary,
        "reuse_decomposition": reuse_decomposition,
        "session_skew": session_skew,
        "append_delta_stats": append_delta_stats,
    }


def build_kv_footprint_summary(input_tokens: list[float], kv_bytes_per_token: float) -> JsonDict:
    if kv_bytes_per_token <= 0:
        return {
            "status": "unavailable",
            "reason": "kv_bytes_per_token was not provided",
            "needed_input": "--kv-bytes-per-token <model-specific bytes per token>",
            "formula": "kv_bytes_per_request = input_tokens * kv_bytes_per_token",
        }
    footprints = [v * kv_bytes_per_token for v in input_tokens]
    return {
        "status": "available",
        "kv_bytes_per_token": kv_bytes_per_token,
        "formula": "kv_bytes_per_request = input_tokens * kv_bytes_per_token",
        "kv_bytes_per_request": stats(footprints),
        "kv_mib_per_request": stats([v / (1024 * 1024) for v in footprints]),
        "total_kv_gib": sum(footprints) / (1024 ** 3),
    }


def build_append_delta_stats(records: list[JsonDict], unmatched_breakdown: list[JsonDict]) -> JsonDict:
    rows = [r for r in records if has_number(r.get("input_tokens")) and has_number(r.get("cached_tokens"))]
    source = "analyzed_records"
    if not rows:
        rows = [
            r for r in unmatched_breakdown
            if has_number(r.get("input_tokens")) and has_number(r.get("cached_tokens"))
        ]
        source = "unmatched_breakdown_without_session_join"
    if not rows:
        return {
            "status": "unavailable",
            "reason": "missing per-request cached token field",
            "needed_fields": ["cached_tokens or cache_hit", "input_length/input_tokens"],
            "formula": "uncached_tokens = input_tokens - cached_tokens",
        }
    uncached = []
    cached = []
    append_ratios = []
    inputs = []
    long_small_append = 0
    for r in rows:
        inp = to_float(r.get("input_tokens"))
        cache = to_float(r.get("cached_tokens"))
        if inp is None or cache is None:
            continue
        unc = max(inp - cache, 0.0)
        inputs.append(inp)
        cached.append(cache)
        uncached.append(unc)
        if inp > 0:
            ratio = unc / inp
            append_ratios.append(ratio)
            if inp >= 32_000 and ratio <= 0.2:
                long_small_append += 1
    return {
        "status": "available",
        "source": source,
        "request_count": len(rows),
        "input_tokens": stats(inputs),
        "cached_tokens": stats(cached),
        "uncached_tokens": stats(uncached),
        "uncached_to_input_ratio": stats(append_ratios),
        "total_input_tokens": sum(inputs),
        "total_cached_tokens": sum(cached),
        "total_uncached_tokens": sum(uncached),
        "overall_cached_fraction": sum(cached) / sum(inputs) if sum(inputs) > 0 else None,
        "long_prompt_small_append_count": long_small_append,
        "long_prompt_small_append_fraction": long_small_append / len(rows) if rows else None,
        "formula": "uncached_tokens = max(input_tokens - cached_tokens, 0)",
    }


def build_session_skew(records: list[JsonDict]) -> JsonDict:
    sessions = group_by_session(records)
    if not sessions:
        return {
            "status": "unavailable",
            "reason": "missing session_id",
        }
    rows: list[JsonDict] = []
    uncached_available = True
    for sid, items in sessions.items():
        input_sum = sum(to_float(r.get("input_tokens")) or 0.0 for r in items)
        output_sum = sum(to_float(r.get("output_tokens")) or 0.0 for r in items)
        cached_values = [to_float(r.get("cached_tokens")) for r in items]
        if any(v is None for v in cached_values):
            uncached_available = False
            uncached_sum = None
        else:
            uncached_sum = sum(
                max((to_float(r.get("input_tokens")) or 0.0) - (to_float(r.get("cached_tokens")) or 0.0), 0.0)
                for r in items
            )
        rows.append({
            "session_id": sid,
            "turns": len(items),
            "input_tokens": input_sum,
            "output_tokens": output_sum,
            "uncached_tokens": uncached_sum,
        })
    input_values = [float(r["input_tokens"]) for r in rows]
    output_values = [float(r["output_tokens"]) for r in rows]
    uncached_values = [float(r["uncached_tokens"]) for r in rows if r["uncached_tokens"] is not None]
    rows_by_input = sorted(rows, key=lambda r: (-float(r["input_tokens"]), str(r["session_id"])))
    return {
        "status": "available",
        "session_count": len(rows),
        "turns_per_session": stats([float(r["turns"]) for r in rows]),
        "cumulative_input_tokens_per_session": stats(input_values),
        "cumulative_output_tokens_per_session": stats(output_values),
        "cumulative_uncached_tokens_per_session": stats(uncached_values) if uncached_available else None,
        "uncached_tokens_status": "available" if uncached_available else "unavailable_missing_cached_tokens",
        "top_session_contribution": {
            "input_tokens": top_contribution(input_values),
            "output_tokens": top_contribution(output_values),
            "uncached_tokens": top_contribution(uncached_values) if uncached_available else None,
        },
        "top_sessions_by_input_tokens": rows_by_input[:20],
        "lorenz_input_tokens": lorenz_points(input_values),
    }


def build_reuse_decomposition(
    *,
    records: list[JsonDict],
    block_size: int,
    shared_prefix_min_sessions: int,
    system_prefix_blocks: int,
) -> JsonDict:
    with_cached = [r for r in records if has_number(r.get("cached_tokens"))]
    with_hash = [r for r in records if r.get("hash_ids")]
    potential = trace_potential_reuse(records)
    if not with_cached:
        return {
            "status": "unavailable",
            "reason": "missing per-request cached token field",
            "needed_fields": ["cached_tokens or cache_hit"],
            "trace_potential_reuse": potential,
        }
    if not with_hash:
        return {
            "status": "unavailable",
            "reason": "missing hash_ids needed to classify cached tokens as intra-session or cross-session",
            "needed_fields": ["hash_ids"],
            "trace_potential_reuse": potential,
        }

    sessions_by_hash: dict[int, set[str]] = collections.defaultdict(set)
    for r in records:
        sid = r.get("session_id")
        if sid is None:
            continue
        for hid in r.get("hash_ids") or []:
            sessions_by_hash[int(hid)].add(str(sid))

    ordered = sorted(
        records,
        key=lambda r: (
            value_or_inf(r.get("scheduled_dispatch_s")),
            value_or_inf(r.get("source_index")),
            str(r.get("request_id", "")),
        ),
    )
    seen_sessions_by_hash: dict[int, set[str]] = collections.defaultdict(set)
    buckets = {
        "intra_session_tokens": 0.0,
        "cross_session_tokens": 0.0,
        "shared_system_prefix_tokens": 0.0,
        "cached_but_unclassified_tokens": 0.0,
    }
    request_rows: list[JsonDict] = []
    for r in ordered:
        sid = str(r.get("session_id")) if r.get("session_id") is not None else None
        hash_ids = [int(h) for h in (r.get("hash_ids") or [])]
        cached_remaining = max(to_float(r.get("cached_tokens")) or 0.0, 0.0)
        local = {key: 0.0 for key in buckets}
        for pos, hid in enumerate(hash_ids):
            if cached_remaining <= 0:
                break
            weight = min(float(block_size), cached_remaining)
            prior_sessions = seen_sessions_by_hash.get(hid, set())
            if (
                len(sessions_by_hash.get(hid, set())) >= shared_prefix_min_sessions
                and pos < system_prefix_blocks
            ):
                key = "shared_system_prefix_tokens"
            elif sid is not None and sid in prior_sessions:
                key = "intra_session_tokens"
            elif prior_sessions:
                key = "cross_session_tokens"
            else:
                key = "cached_but_unclassified_tokens"
            buckets[key] += weight
            local[key] += weight
            cached_remaining -= weight
        if any(v > 0 for v in local.values()):
            request_rows.append({
                "request_id": r.get("request_id"),
                "session_id": sid,
                **local,
            })
        if sid is not None:
            for hid in hash_ids:
                seen_sessions_by_hash[int(hid)].add(sid)

    total = sum(buckets.values())
    fractions = {
        key.replace("_tokens", "_fraction"): value / total if total > 0 else None
        for key, value in buckets.items()
    }
    return {
        "status": "available" if total > 0 else "available_no_cached_reuse_observed",
        "block_size": block_size,
        "shared_prefix_min_sessions": shared_prefix_min_sessions,
        "system_prefix_blocks": system_prefix_blocks,
        "method": (
            "Classify each cached hash block by earlier occurrences. Blocks in the first "
            "system_prefix_blocks positions that appear in at least shared_prefix_min_sessions "
            "sessions are counted as shared/system-prefix."
        ),
        **buckets,
        **fractions,
        "total_classified_cached_tokens": total,
        "request_level_reuse_rows": request_rows[:200],
        "trace_potential_reuse": potential,
    }


def trace_potential_reuse(records: list[JsonDict]) -> JsonDict:
    seen_sessions_by_hash: dict[int, set[str]] = collections.defaultdict(set)
    counts = {
        "first_seen_block_occurrences": 0,
        "intra_session_reused_block_occurrences": 0,
        "cross_session_reused_block_occurrences": 0,
    }
    ordered = sorted(
        records,
        key=lambda r: (
            value_or_inf(r.get("scheduled_dispatch_s")),
            value_or_inf(r.get("source_index")),
            str(r.get("request_id", "")),
        ),
    )
    for r in ordered:
        sid = str(r.get("session_id")) if r.get("session_id") is not None else None
        if sid is None:
            continue
        for hid in r.get("hash_ids") or []:
            hid = int(hid)
            prior_sessions = seen_sessions_by_hash.get(hid, set())
            if sid in prior_sessions:
                counts["intra_session_reused_block_occurrences"] += 1
            elif prior_sessions:
                counts["cross_session_reused_block_occurrences"] += 1
            else:
                counts["first_seen_block_occurrences"] += 1
            seen_sessions_by_hash[hid].add(sid)
    total_reused = (
        counts["intra_session_reused_block_occurrences"]
        + counts["cross_session_reused_block_occurrences"]
    )
    return {
        **counts,
        "intra_session_reuse_fraction": counts["intra_session_reused_block_occurrences"] / total_reused if total_reused else None,
        "cross_session_reuse_fraction": counts["cross_session_reused_block_occurrences"] / total_reused if total_reused else None,
        "note": "Potential reuse from trace hash_ids only; it is not actual cache-hit reuse.",
    }


def generate_figures(
    out_dir: Path,
    records: list[JsonDict],
    batch0: JsonDict,
    batch1: JsonDict,
    no_figures: bool,
) -> JsonDict:
    if no_figures:
        return {"status": "skipped", "reason": "--no-figures"}
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except Exception as exc:  # pragma: no cover - depends on optional dependency.
        return {
            "status": "skipped",
            "reason": f"matplotlib unavailable: {exc!r}",
        }

    figures_dir = out_dir / "figures"
    created: list[str] = []
    skipped: dict[str, str] = {}

    def save_current(name: str) -> None:
        for suffix in ["png", "pdf"]:
            path = figures_dir / f"{name}.{suffix}"
            plt.savefig(path, bbox_inches="tight")
            created.append(str(path.relative_to(out_dir)))
        plt.close()

    def cdf_plot(values: list[float], name: str, xlabel: str, label: str | None = None) -> None:
        clean = sorted(v for v in values if v is not None and math.isfinite(v))
        if not clean:
            skipped[name] = "no data"
            return
        y = [(i + 1) / len(clean) for i in range(len(clean))]
        plt.figure(figsize=(6, 4))
        plt.plot(clean, y, label=label or xlabel)
        if label:
            plt.legend()
        plt.xlabel(xlabel)
        plt.ylabel("CDF")
        plt.grid(True, alpha=0.3)
        save_current(name)

    starts = batch0["session_arrival_stats"].get("session_start_offsets_s") or []
    cdf_plot(starts, "session_start_time_cdf", "session start offset (s)")

    actual_per_session = batch0["session_concurrency"].get("per_session") or []
    max_inflight = [float(r["max_inflight"]) for r in actual_per_session if "max_inflight" in r]
    if max_inflight:
        plt.figure(figsize=(6, 4))
        bins = range(1, int(max(max_inflight)) + 2)
        plt.hist(max_inflight, bins=bins, align="left", rwidth=0.8)
        plt.xlabel("per-session max in-flight")
        plt.ylabel("sessions")
        plt.grid(True, alpha=0.3)
        save_current("per_session_max_inflight_hist")
    else:
        skipped["per_session_max_inflight_hist"] = "actual interval data unavailable"

    turns = []
    sessions = group_by_session(records)
    for rows in sessions.values():
        turns.append(float(len(rows)))
    cdf_plot(turns, "turns_per_session_cdf", "turns per session")

    turn_values = batch0["turn_interval_stats"].get("scheduled_turn_interval_values_s") or []
    cdf_plot(turn_values, "turn_interarrival_cdf", "scheduled turn inter-arrival (s)")

    input_values = clean_numbers(r.get("input_tokens") for r in records)
    output_values = clean_numbers(r.get("output_tokens") for r in records)
    if input_values or output_values:
        plt.figure(figsize=(6, 4))
        if input_values:
            xs = sorted(input_values)
            ys = [(i + 1) / len(xs) for i in range(len(xs))]
            plt.plot(xs, ys, label="input")
        if output_values:
            xs = sorted(output_values)
            ys = [(i + 1) / len(xs) for i in range(len(xs))]
            plt.plot(xs, ys, label="output")
        plt.xlabel("tokens")
        plt.ylabel("CDF")
        plt.legend()
        plt.grid(True, alpha=0.3)
        save_current("input_output_token_cdf")
    else:
        skipped["input_output_token_cdf"] = "no token data"

    ratios = []
    for r in records:
        inp = to_float(r.get("input_tokens"))
        out = to_float(r.get("output_tokens"))
        if inp is not None and out is not None and out > 0:
            ratios.append(inp / out)
    cdf_plot(ratios, "input_output_ratio_cdf", "input/output token ratio")

    kv_summary = batch1["kv_footprint_summary"]
    if kv_summary.get("status") == "available":
        kv_bpt = float(kv_summary["kv_bytes_per_token"])
        cdf_plot([v * kv_bpt / (1024 * 1024) for v in input_values], "kv_footprint_cdf", "KV footprint (MiB)")
    else:
        skipped["kv_footprint_cdf"] = kv_summary.get("reason", "unavailable")

    reuse = batch1["reuse_decomposition"]
    if reuse.get("status", "").startswith("available"):
        labels = [
            "intra",
            "cross",
            "shared/system",
            "unclassified",
        ]
        values = [
            reuse.get("intra_session_tokens") or 0,
            reuse.get("cross_session_tokens") or 0,
            reuse.get("shared_system_prefix_tokens") or 0,
            reuse.get("cached_but_unclassified_tokens") or 0,
        ]
        plt.figure(figsize=(6, 2.5))
        left = 0.0
        for label, value in zip(labels, values):
            plt.barh(["cached reuse"], [value], left=left, label=label)
            left += value
        plt.xlabel("tokens")
        plt.legend(loc="best", fontsize=8)
        save_current("reuse_decomposition_stacked_bar")
    else:
        skipped["reuse_decomposition_stacked_bar"] = reuse.get("reason", "unavailable")

    skew = batch1["session_skew"]
    lorenz = skew.get("lorenz_input_tokens") or []
    if lorenz:
        plt.figure(figsize=(5, 5))
        plt.plot([p["session_fraction"] for p in lorenz], [p["token_fraction"] for p in lorenz])
        plt.plot([0, 1], [0, 1], linestyle="--", color="gray", linewidth=1)
        plt.xlabel("fraction of sessions")
        plt.ylabel("fraction of input tokens")
        plt.grid(True, alpha=0.3)
        save_current("per_session_token_mass_lorenz")
    else:
        skipped["per_session_token_mass_lorenz"] = "no session skew data"

    top_sessions = skew.get("top_sessions_by_input_tokens") or []
    if top_sessions:
        top = top_sessions[:10]
        plt.figure(figsize=(7, 4))
        plt.bar([str(r["session_id"]) for r in top], [float(r["input_tokens"]) for r in top])
        plt.xlabel("session id")
        plt.ylabel("input tokens")
        plt.xticks(rotation=45, ha="right", fontsize=8)
        save_current("top_sessions_token_contribution")
    else:
        skipped["top_sessions_token_contribution"] = "no session skew data"

    append_rows = [r for r in records if has_number(r.get("input_tokens")) and has_number(r.get("cached_tokens"))]
    if append_rows:
        x = [float(r["input_tokens"]) for r in append_rows]
        y = [max(float(r["input_tokens"]) - float(r["cached_tokens"]), 0.0) for r in append_rows]
        plt.figure(figsize=(6, 4))
        plt.scatter(x, y, s=10, alpha=0.5)
        plt.xlabel("input tokens")
        plt.ylabel("uncached tokens")
        plt.grid(True, alpha=0.3)
        save_current("total_input_vs_uncached_scatter")
    else:
        skipped["total_input_vs_uncached_scatter"] = "missing cached_tokens/cache_hit"

    return {
        "status": "available",
        "created": created,
        "skipped": skipped,
    }


def build_manifest(
    *,
    args: argparse.Namespace,
    config: JsonDict,
    out_dir: Path,
    started_at: dt.datetime,
    finished_at: dt.datetime,
    batch0: JsonDict,
) -> JsonDict:
    time_scale = args.time_scale
    if time_scale is None:
        time_scale = to_float(config.get("time_scale"))
    request_limit = args.request_limit
    if request_limit is None:
        request_limit = to_int(first_present(config, ["requests", "request_limit"]))
    policy = args.policy or str(config.get("policy", ""))
    trace_info = file_info(args.trace, hash_inputs=args.hash_inputs)
    return {
        "git_commit": git_commit(),
        "host": platform.node(),
        "gpu_type": args.gpu_type,
        "gpu_count": args.gpu_count,
        "canonical_trace_data_sources": {
            "dash0_formatted_trace_dir": "~/ali-trace/trace-glm5.1-formatted/",
            "dash0_raw_trace_dir": "~/ali-trace/trace-glm5.1/",
            "usage_note": (
                "Full trace analysis can be run CPU-only on dash0, or the needed "
                "JSONL files can be copied/rsynced from dash0 to this machine before "
                "running this analyzer."
            ),
        },
        "input_requirements": {
            "trace_jsonl": [
                "chat_id",
                "parent_chat_id",
                "timestamp",
                "input_length",
                "output_length",
                "turn",
                "hash_ids",
                "optional session_id",
            ],
            "metrics_jsonl": [
                "request_id",
                "session_id",
                "trace_timestamp_s",
                "input_length",
                "output_length",
                "latency_s",
                "ttft_s",
                "tpot_s",
                "error",
                "optional cached_tokens",
            ],
            "actual_sequentiality_proof": [
                "per-request dispatch timestamp",
                "per-request finish or error/timeout timestamp",
                "request_id join to trace/metrics when timing source is separate",
            ],
            "reuse_decomposition": [
                "cached_tokens or cache_hit",
                "hash_ids",
                "session_id",
            ],
        },
        "trace_path": str(args.trace) if args.trace else "",
        "trace_sha256": trace_info.get("sha256", ""),
        "trace_file_info": trace_info,
        "policy": policy,
        "launch_command": args.launch_command or " ".join(sys.argv),
        "request_limit": request_limit,
        "time_scale": time_scale,
        "session_sampling_method": args.session_sampling_method,
        "session_sequential": batch0["session_concurrency"].get("session_sequential"),
        "start_time": started_at.isoformat(),
        "end_time": finished_at.isoformat(),
        "output_dir": str(out_dir),
    }


def render_invalid_runs(batch0: JsonDict) -> str:
    cls = batch0["classification"]
    conc = batch0["session_concurrency"]
    lines = [
        "# Invalid Runs",
        "",
        f"Classification: `{cls['label']}`",
        "",
        f"Reason: {cls['reason']}",
        "",
        "## Sequentiality",
        "",
        f"- Actual interval status: `{conc['status']}`",
        f"- Session sequential: `{conc['session_sequential']}`",
        f"- Max in-flight per session: `{conc['max_inflight_per_session']}`",
        f"- Sessions with overlap: `{conc['sessions_with_overlap']}`",
        "",
    ]
    if conc.get("overlap_examples"):
        lines.append("## Overlap Examples")
        lines.append("")
        for item in conc["overlap_examples"][:10]:
            lines.append(
                f"- session `{item['session_id']}` request `{item['request_id']}` "
                f"started {item['overlap_s']:.6g}s before `{item['previous_request_id']}` finished"
            )
        lines.append("")
    if cls["label"] == "invalid_for_online_claim":
        lines.extend([
            "## Claims Still Supported",
            "",
            "- Workload shape and trace-level token/session distributions.",
            "- Offline or stress-test observations that do not require online per-session sequentiality.",
            "",
            "## Claims Not Supported",
            "",
            "- Online-serving SRR/SLO claims that require at most one in-flight turn per session.",
            "",
        ])
    return "\n".join(lines)


def render_summary_md(summary: JsonDict, batch0: JsonDict, batch1: JsonDict) -> str:
    prof = batch0["trace_profile"]
    conc = batch0["session_concurrency"]
    workload = batch1["workload_summary"]
    kv = batch1["kv_footprint_summary"]
    reuse = batch1["reuse_decomposition"]
    append = batch1["append_delta_stats"]
    lines = [
        "# Characterization Summary",
        "",
        f"- Classification: `{summary['classification']['label']}`",
        f"- Analyzed records: {summary['analyzed_records']}",
        f"- Sessions: {prof.get('session_count')}",
        f"- Attempted/completed/errors: {prof.get('attempted_requests')} / {prof.get('completed_requests')} / {prof.get('error_requests')}",
        f"- Goodput: {fmt(prof.get('goodput_per_s'))} req/s",
        "",
        "## Batch 0",
        "",
        f"- Actual interval status: `{conc['status']}`",
        f"- Session sequential: `{conc['session_sequential']}`",
        f"- Max in-flight per session: `{conc['max_inflight_per_session']}`",
        "",
        "## Batch 1",
        "",
        f"- Input tokens p50/p90/p99: {stats_line(workload.get('input_tokens'))}",
        f"- Output tokens p50/p90/p99: {stats_line(workload.get('output_tokens'))}",
        f"- KV footprint: `{kv.get('status')}`",
        f"- Reuse decomposition: `{reuse.get('status')}`",
        f"- Append/uncached stats: `{append.get('status')}`",
        "",
        "Raw data and figures are in this output directory. See `audit.md` for unavailable fields and claim checks.",
        "",
    ]
    return "\n".join(lines)


def render_audit_md(batch0: JsonDict, batch1: JsonDict) -> str:
    conc = batch0["session_concurrency"]
    prof = batch0["trace_profile"]
    workload = batch1["workload_summary"]
    kv = batch1["kv_footprint_summary"]
    reuse = batch1["reuse_decomposition"]
    skew = batch1["session_skew"]
    append = batch1["append_delta_stats"]
    top = skew.get("top_session_contribution") or {}
    input_top = top.get("input_tokens") or {}
    lines = [
        "# Audit",
        "",
        "## Batch 0 Checks",
        "",
        f"1. Does the main trace satisfy `max_inflight_per_session == 1`? {answer_sequential(conc)}",
        f"2. If not, is the run labeled stress or invalid? Classification is `{batch0['classification']['label']}`.",
        f"3. Are attempted/completed/error counts included? Attempted={prof.get('attempted_requests')}, completed={prof.get('completed_requests')}, errors={prof.get('error_requests')}.",
        f"4. Are latency percentiles over successes, with goodput reported? Success-only={prof.get('latency_percentiles_over_successes_only')}, goodput={fmt(prof.get('goodput_per_s'))} req/s.",
        "",
        "## Batch 1 Checks",
        "",
        f"1. Input p50/p90/p99: {stats_line(workload.get('input_tokens'))}.",
        f"2. Output p50/p90/p99: {stats_line(workload.get('output_tokens'))}.",
        f"3. Estimated KV footprint p50/p90/p99: {kv_line(kv)}.",
        f"4. Fraction of reuse intra-session: {reuse_fraction_line(reuse, 'intra_session_fraction')}.",
        f"5. Top 1% / 5% session token mass: {fmt(input_top.get('top_1pct_fraction'))} / {fmt(input_top.get('top_5pct_fraction'))}.",
        f"6. Are long prompts often small appends after cache reuse? {append_line(append)}",
        "",
        "## Missing Inputs",
        "",
    ]
    missing = collect_missing_inputs(conc, kv, reuse, append)
    if missing:
        lines.extend(f"- {item}" for item in missing)
    else:
        lines.append("- None detected by the analyzer.")
    lines.append("")
    return "\n".join(lines)


def answer_sequential(conc: JsonDict) -> str:
    if conc.get("session_sequential") is True:
        return "yes, actual interval data shows max in-flight <= 1."
    if conc.get("session_sequential") is False:
        return "no, overlap was observed."
    return (
        "unavailable, because actual dispatch and finish/error timestamps "
        "are missing or only partially available."
    )


def collect_missing_inputs(conc: JsonDict, kv: JsonDict, reuse: JsonDict, append: JsonDict) -> list[str]:
    missing: list[str] = []
    if conc.get("status") != "available":
        missing.append(
            "Batch 0 actual sequentiality needs per-request dispatch timestamp plus finish/error timestamp."
        )
    if kv.get("status") != "available":
        missing.append(str(kv.get("needed_input", "KV bytes per token is required.")))
    if reuse.get("status") == "unavailable":
        missing.append("Reuse decomposition needs `cached_tokens` or `cache_hit` and `hash_ids`.")
    if append.get("status") == "unavailable":
        missing.append("Append stats need `cached_tokens` or `cache_hit` plus input length.")
    return missing


def relative_output_list(out_dir: Path) -> list[str]:
    return sorted(
        str(path.relative_to(out_dir))
        for path in out_dir.rglob("*")
        if path.is_file()
    )


def group_by_session(records: list[JsonDict]) -> dict[str, list[JsonDict]]:
    sessions: dict[str, list[JsonDict]] = collections.defaultdict(list)
    for r in records:
        sid = r.get("session_id")
        if sid is not None:
            sessions[str(sid)].append(r)
    return dict(sessions)


def stats(values: list[float] | list[int]) -> JsonDict | None:
    clean = sorted(float(v) for v in values if v is not None and math.isfinite(float(v)))
    if not clean:
        return None
    return {
        "count": len(clean),
        "mean": statistics.fmean(clean),
        "min": clean[0],
        "p50": percentile(clean, 0.50),
        "p90": percentile(clean, 0.90),
        "p95": percentile(clean, 0.95),
        "p99": percentile(clean, 0.99),
        "max": clean[-1],
    }


def percentile(sorted_values: list[float], pct: float) -> float:
    if not sorted_values:
        raise ValueError("percentile requires at least one value")
    if len(sorted_values) == 1:
        return sorted_values[0]
    rank = pct * (len(sorted_values) - 1)
    lo = int(rank)
    hi = min(lo + 1, len(sorted_values) - 1)
    frac = rank - lo
    return sorted_values[lo] * (1 - frac) + sorted_values[hi] * frac


def clean_numbers(values: Any) -> list[float]:
    out: list[float] = []
    for value in values:
        num = to_float(value)
        if num is not None and math.isfinite(num):
            out.append(num)
    return out


def top_contribution(values: list[float]) -> JsonDict:
    clean = sorted([v for v in values if v is not None and math.isfinite(v)], reverse=True)
    total = sum(clean)
    if not clean or total <= 0:
        return {
            "top_1pct_fraction": None,
            "top_5pct_fraction": None,
            "top_10pct_fraction": None,
            "top_20pct_fraction": None,
        }
    def frac(pct: float) -> float:
        k = max(1, math.ceil(len(clean) * pct))
        return sum(clean[:k]) / total
    return {
        "top_1pct_fraction": frac(0.01),
        "top_5pct_fraction": frac(0.05),
        "top_10pct_fraction": frac(0.10),
        "top_20pct_fraction": frac(0.20),
    }


def lorenz_points(values: list[float]) -> list[JsonDict]:
    clean = sorted(v for v in values if v is not None and math.isfinite(v) and v >= 0)
    total = sum(clean)
    if not clean or total <= 0:
        return []
    points = [{"session_fraction": 0.0, "token_fraction": 0.0}]
    cumulative = 0.0
    for idx, value in enumerate(clean, start=1):
        cumulative += value
        points.append({
            "session_fraction": idx / len(clean),
            "token_fraction": cumulative / total,
        })
    return points


def count_gt(values: list[float], threshold: float) -> JsonDict:
    if not values:
        return {"count": 0, "fraction": None}
    count = sum(1 for v in values if v > threshold)
    return {"count": count, "fraction": count / len(values)}


def first_present(row: JsonDict, keys: list[str], default: Any = None) -> Any:
    for key in keys:
        if key in row and row[key] is not None:
            return row[key]
    return default


def to_float(value: Any, default: float | None = None) -> float | None:
    if value is None:
        return default
    try:
        return float(value)
    except (TypeError, ValueError):
        return default


def to_int(value: Any, default: int | None = None) -> int | None:
    if value is None:
        return default
    try:
        return int(value)
    except (TypeError, ValueError):
        return default


def has_number(value: Any) -> bool:
    return to_float(value) is not None


def value_or_inf(value: Any) -> float:
    num = to_float(value)
    return num if num is not None else float("inf")


def fmt(value: Any) -> str:
    num = to_float(value)
    if num is None:
        return "unavailable"
    if abs(num) >= 1000:
        return f"{num:,.3f}"
    return f"{num:.6g}"


def stats_line(s: JsonDict | None) -> str:
    if not s:
        return "unavailable"
    return f"{fmt(s.get('p50'))} / {fmt(s.get('p90'))} / {fmt(s.get('p99'))}"


def kv_line(kv: JsonDict) -> str:
    if kv.get("status") != "available":
        return f"unavailable ({kv.get('reason')})"
    return stats_line(kv.get("kv_mib_per_request")) + " MiB"


def reuse_fraction_line(reuse: JsonDict, key: str) -> str:
    if not reuse.get("status", "").startswith("available"):
        return f"unavailable ({reuse.get('reason')})"
    return fmt(reuse.get(key))


def append_line(append: JsonDict) -> str:
    if append.get("status") != "available":
        return f"unavailable ({append.get('reason')})"
    return (
        f"long_prompt_small_append_fraction={fmt(append.get('long_prompt_small_append_fraction'))}; "
        f"uncached/input p50/p90/p99={stats_line(append.get('uncached_to_input_ratio'))}"
    )


def write_json(path: Path, data: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def write_jsonl(path: Path, rows: list[JsonDict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, sort_keys=True) + "\n")


def sha256_file(path: Path | None) -> str:
    if path is None or not path.exists():
        return ""
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def file_info(path: Path | None, *, hash_inputs: bool) -> JsonDict:
    if path is None:
        return {
            "path": "",
            "exists": False,
            "sha256_status": "unavailable_no_path",
        }
    info: JsonDict = {
        "path": str(path),
        "exists": path.exists(),
    }
    if not path.exists():
        info["sha256_status"] = "unavailable_missing_file"
        return info
    stat = path.stat()
    info.update({
        "size_bytes": stat.st_size,
        "mtime_s": stat.st_mtime,
    })
    if hash_inputs:
        info["sha256"] = sha256_file(path)
        info["sha256_status"] = "available"
    else:
        info["sha256"] = ""
        info["sha256_status"] = "skipped_use_--hash-inputs"
    return info


def git_commit() -> str:
    try:
        result = subprocess.run(
            ["git", "rev-parse", "HEAD"],
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            text=True,
        )
    except Exception:
        return ""
    return result.stdout.strip()


if __name__ == "__main__":
    main()
