import argparse
import csv
import json
import statistics
import time
from collections import defaultdict
from pathlib import Path
from typing import Any


MEDIAN_FIELDS = (
    "load_seconds",
    "generate_seconds",
    "request_count",
    "output_tokens",
    "output_tokens_per_second",
    "requests_per_second",
    "latency_p50_seconds",
    "latency_p95_seconds",
    "process_rss_before_mib",
    "process_rss_peak_mib",
    "process_rss_after_mib",
    "npu_hbm_before_mib",
    "npu_hbm_peak_mib",
    "npu_hbm_after_mib",
)


def percent_change(value: float | None,
                   baseline: float | None) -> float | None:
    if value is None or baseline is None or baseline == 0:
        return None
    return (value - baseline) / baseline * 100.0


def hbm_saved_gib(baseline_mib: float | None,
                  value_mib: float | None) -> float | None:
    if baseline_mib is None or value_mib is None:
        return None
    return (baseline_mib - value_mib) / 1024.0


def median_or_none(records: list[dict[str, Any]], field: str) -> float | None:
    values = [record[field] for record in records
              if isinstance(record.get(field), (int, float))]
    return statistics.median(values) if values else None


def aggregate_records(records: list[dict[str, Any]],
                      baseline_label: str) -> list[dict[str, Any]]:
    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for record in records:
        label = record.get("label")
        if not isinstance(label, str) or not label:
            raise ValueError("every benchmark record must have a non-empty label")
        grouped[label].append(record)
    if baseline_label not in grouped:
        raise ValueError(f"baseline label {baseline_label!r} is missing")

    rows: list[dict[str, Any]] = []
    for label in sorted(grouped, key=lambda item: (item != baseline_label, item)):
        case_records = grouped[label]
        first = case_records[0]
        row: dict[str, Any] = {
            "label": label,
            "repeats": len(case_records),
            "model": first.get("model"),
            "offload_group_size": first.get("offload_group_size"),
            "offload_num_in_group": first.get("offload_num_in_group"),
        }
        for field in MEDIAN_FIELDS:
            row[field] = median_or_none(case_records, field)
        rows.append(row)

    baseline = next(row for row in rows if row["label"] == baseline_label)
    for row in rows:
        if row["label"] == baseline_label:
            row["throughput_change_percent"] = 0.0
            row["hbm_saved_gib"] = 0.0
        else:
            row["throughput_change_percent"] = percent_change(
                row["output_tokens_per_second"],
                baseline["output_tokens_per_second"],
            )
            row["hbm_saved_gib"] = hbm_saved_gib(
                baseline["npu_hbm_peak_mib"], row["npu_hbm_peak_mib"]
            )
    return rows


def load_jsonl(path: Path) -> list[dict[str, Any]]:
    records = []
    with path.open(encoding="utf-8") as handle:
        for line_number, line in enumerate(handle, 1):
            if not line.strip():
                continue
            try:
                record = json.loads(line)
            except json.JSONDecodeError as exc:
                raise ValueError(f"invalid JSON on {path}:{line_number}: {exc}") from exc
            if not isinstance(record, dict):
                raise ValueError(f"record on {path}:{line_number} is not an object")
            records.append(record)
    if not records:
        raise ValueError(f"no benchmark records found in {path}")
    return records


def format_cell(value: Any) -> str:
    if value is None:
        return "N/A"
    if isinstance(value, float):
        return f"{value:.4f}"
    return str(value)


def write_outputs(rows: list[dict[str, Any]], output_dir: Path,
                  baseline_label: str, input_path: Path) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    fieldnames = list(rows[0])
    csv_path = output_dir / "performance-summary.csv"
    with csv_path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    markdown_fields = (
        "label", "repeats", "output_tokens_per_second", "requests_per_second",
        "latency_p50_seconds", "latency_p95_seconds", "npu_hbm_peak_mib",
        "throughput_change_percent", "hbm_saved_gib",
    )
    lines = [
        "| " + " | ".join(markdown_fields) + " |",
        "| " + " | ".join("---" for _ in markdown_fields) + " |",
    ]
    lines.extend(
        "| " + " | ".join(format_cell(row[field]) for field in markdown_fields) + " |"
        for row in rows
    )
    (output_dir / "performance-summary.md").write_text(
        "\n".join(lines) + "\n", encoding="utf-8"
    )

    payload = {
        "generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "baseline_label": baseline_label,
        "input_jsonl": str(input_path.resolve()),
        "aggregation": "median",
        "rows": rows,
    }
    (output_dir / "performance-summary.json").write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--baseline-label", default="baseline_0_1")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    rows = aggregate_records(load_jsonl(args.input), args.baseline_label)
    write_outputs(rows, args.output_dir, args.baseline_label, args.input)
    return 0


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