#!/usr/bin/env python3
"""Low-overhead collectors and result summarizer for Kimi-K3 PD attribution."""

from __future__ import annotations

import argparse
import csv
import json
import re
import statistics
import subprocess
import time
from pathlib import Path


HCA_COUNTERS = (
    "port_xmit_data",
    "port_rcv_data",
    "port_xmit_wait",
    "port_rcv_errors",
    "port_xmit_discards",
)


def read_int(path: Path) -> int | None:
    try:
        return int(path.read_text().strip())
    except (FileNotFoundError, PermissionError, ValueError):
        return None


def read_hcas() -> dict[str, dict[str, int | None]]:
    result = {}
    for index in range(4):
        hca = f"mlx5_{index}"
        root = Path(f"/sys/class/infiniband/{hca}/ports/1/counters")
        result[hca] = {name: read_int(root / name) for name in HCA_COUNTERS}
    return result


def read_gpus() -> list[dict[str, float | int | str | None]]:
    fields = (
        "index,timestamp,utilization.gpu,memory.used,power.draw,clocks.sm,"
        "pcie.link.gen.current,pcie.link.width.current"
    )
    proc = subprocess.run(
        [
            "nvidia-smi",
            f"--query-gpu={fields}",
            "--format=csv,noheader,nounits",
        ],
        check=False,
        capture_output=True,
        text=True,
    )
    rows = []
    for line in proc.stdout.splitlines():
        values = [value.strip() for value in line.split(",")]
        if len(values) != 8:
            continue
        parsed: list[float | int | str | None] = [int(values[0]), values[1]]
        for value in values[2:]:
            try:
                parsed.append(float(value))
            except ValueError:
                parsed.append(None)
        rows.append(
            dict(
                zip(
                    (
                        "index",
                        "timestamp",
                        "gpu_util_pct",
                        "memory_used_mib",
                        "power_w",
                        "sm_clock_mhz",
                        "pcie_gen",
                        "pcie_width",
                    ),
                    parsed,
                )
            )
        )
    return rows


def sample(args: argparse.Namespace) -> None:
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    deadline = time.monotonic() + args.duration
    with out.open("w") as handle:
        while time.monotonic() < deadline:
            started = time.monotonic()
            record = {
                "wall_time": time.time(),
                "monotonic": started,
                "hcas": read_hcas(),
                "gpus": read_gpus(),
            }
            handle.write(json.dumps(record, separators=(",", ":")) + "\n")
            handle.flush()
            time.sleep(max(0.0, args.interval - (time.monotonic() - started)))


def percentile(values: list[float], q: float) -> float | None:
    if not values:
        return None
    ordered = sorted(values)
    index = (len(ordered) - 1) * q
    lower = int(index)
    upper = min(lower + 1, len(ordered) - 1)
    fraction = index - lower
    return ordered[lower] * (1 - fraction) + ordered[upper] * fraction


def monitor_summary(path: Path) -> dict:
    rows = [json.loads(line) for line in path.read_text().splitlines() if line]
    hca_rates: dict[str, dict[str, list[float]]] = {}
    gpu_util, gpu_power = [], []
    for previous, current in zip(rows, rows[1:]):
        elapsed = current["monotonic"] - previous["monotonic"]
        if elapsed <= 0:
            continue
        for hca, counters in current["hcas"].items():
            target = hca_rates.setdefault(hca, {"tx_gbps": [], "rx_gbps": []})
            old = previous["hcas"].get(hca, {})
            for counter, output in (
                ("port_xmit_data", "tx_gbps"),
                ("port_rcv_data", "rx_gbps"),
            ):
                before, after = old.get(counter), counters.get(counter)
                if before is not None and after is not None and after >= before:
                    # IB port_*_data counters are in 32-bit words.
                    target[output].append((after - before) * 4 * 8 / elapsed / 1e9)
        for gpu in current.get("gpus", []):
            if gpu.get("gpu_util_pct") is not None:
                gpu_util.append(float(gpu["gpu_util_pct"]))
            if gpu.get("power_w") is not None:
                gpu_power.append(float(gpu["power_w"]))

    hcas = {}
    for hca, rates in hca_rates.items():
        hcas[hca] = {
            name: {
                "mean": statistics.fmean(values) if values else None,
                "p95": percentile(values, 0.95),
                "max": max(values) if values else None,
            }
            for name, values in rates.items()
        }
    return {
        "samples": len(rows),
        "hcas": hcas,
        "gpu_util_pct_mean": statistics.fmean(gpu_util) if gpu_util else None,
        "gpu_util_pct_p95": percentile(gpu_util, 0.95),
        "gpu_power_w_mean": statistics.fmean(gpu_power) if gpu_power else None,
    }


REQ_RE = re.compile(r"ReqTimeStats\(.*?type=(prefill|decode)\): (.*)$")
VALUE_RE = re.compile(r"([a-z_]+)=([0-9.]+)(ms| GB/s| MB)?")


def request_stats(paths: list[Path]) -> dict:
    by_role: dict[str, dict[str, list[float]]] = {"prefill": {}, "decode": {}}
    for path in paths:
        for line in path.read_text(errors="replace").splitlines():
            match = REQ_RE.search(line)
            if not match:
                continue
            role, payload = match.groups()
            for key, value, unit in VALUE_RE.findall(payload):
                number = float(value)
                if unit == "ms":
                    number /= 1000
                by_role[role].setdefault(key, []).append(number)
    result = {}
    for role, fields in by_role.items():
        result[role] = {
            key: {
                "count": len(values),
                "mean_s": statistics.fmean(values),
                "p50_s": percentile(values, 0.50),
                "p95_s": percentile(values, 0.95),
            }
            for key, values in fields.items()
        }
    return result


def benchmark_summary(path: Path) -> dict:
    payload = json.loads(path.read_text().splitlines()[0])
    keys = (
        "duration",
        "completed",
        "total_input_tokens",
        "total_output_tokens",
        "input_throughput",
        "output_throughput",
        "mean_ttft_ms",
        "median_ttft_ms",
        "p95_ttft_ms",
        "mean_tpot_ms",
        "median_tpot_ms",
        "p95_tpot_ms",
        "mean_e2e_latency_ms",
        "p95_e2e_latency_ms",
        "accept_length",
    )
    return {key: payload.get(key) for key in keys}


def analyze(args: argparse.Namespace) -> None:
    root = Path(args.result_root)
    bench_files = sorted((root / "bench").glob("*.jsonl"))
    service_logs = sorted((root / "logs").glob("[pd]_*.log"))
    monitor_files = sorted((root / "telemetry").glob("*.jsonl"))
    summary = {
        "run_id": root.name,
        "benchmarks": {path.name: benchmark_summary(path) for path in bench_files},
        "request_stages": request_stats(service_logs),
        "nodes": {path.stem: monitor_summary(path) for path in monitor_files},
    }
    (root / "summary.json").write_text(json.dumps(summary, indent=2) + "\n")

    rows = []
    for name, item in summary["benchmarks"].items():
        rows.append(
            {
                "file": name,
                "input_tps": item.get("input_throughput"),
                "output_tps": item.get("output_throughput"),
                "ttft_mean_ms": item.get("mean_ttft_ms"),
                "ttft_p95_ms": item.get("p95_ttft_ms"),
                "tpot_mean_ms": item.get("mean_tpot_ms"),
                "e2e_mean_ms": item.get("mean_e2e_latency_ms"),
                "accept_length": item.get("accept_length"),
            }
        )
    with (root / "summary.csv").open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=rows[0].keys() if rows else ["file"])
        writer.writeheader()
        writer.writerows(rows)
    print(json.dumps(summary, indent=2))


def main() -> None:
    parser = argparse.ArgumentParser()
    sub = parser.add_subparsers(dest="command", required=True)
    collect = sub.add_parser("sample")
    collect.add_argument("--out", required=True)
    collect.add_argument("--interval", type=float, default=1.0)
    collect.add_argument("--duration", type=float, default=900.0)
    collect.set_defaults(func=sample)
    report = sub.add_parser("analyze")
    report.add_argument("--result-root", required=True)
    report.set_defaults(func=analyze)
    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()