#!/usr/bin/env python3
"""Extract directly measurable engine/simulator state into one schema.

The schema deliberately keeps common, engine-only, and simulator-only fields
separate.  Missing simulator mechanisms must remain missing; they must not be
filled by a heuristic bottleneck label.
"""

from __future__ import annotations

import csv
import json
import math
from pathlib import Path
from statistics import fmean
from typing import Any, Iterable, Mapping, Sequence


SCHEMA = "telemetry-common-state-v1"


def numeric(values: Iterable[float | int]) -> dict[str, Any]:
    finite = [float(value) for value in values]
    if not finite:
        raise ValueError("numeric summary requires at least one value")
    if any(not math.isfinite(value) for value in finite):
        raise ValueError("numeric summary received a non-finite value")
    mean = fmean(finite)
    variance = fmean((value - mean) ** 2 for value in finite)
    return {
        "n": len(finite),
        "min": min(finite),
        "max": max(finite),
        "mean": mean,
        "cv": math.sqrt(variance) / abs(mean) if mean else 0.0,
        "distinct_n": len(set(finite)),
    }


def load_jsonl(path: Path) -> list[dict[str, Any]]:
    rows = []
    with path.open(encoding="utf-8") as source:
        for line_number, line in enumerate(source, start=1):
            if not line.strip():
                continue
            row = json.loads(line)
            if not isinstance(row, dict):
                raise ValueError(f"{path}:{line_number}: expected a JSON object")
            rows.append(row)
    if not rows:
        raise ValueError(f"{path}: no JSONL records")
    return rows


def _as_number(value: Any, *, name: str) -> float:
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise ValueError(f"{name} must be numeric, got {value!r}")
    result = float(value)
    if not math.isfinite(result):
        raise ValueError(f"{name} must be finite, got {value!r}")
    return result


def _time_weighted_mean(
    records: Sequence[Mapping[str, Any]],
    *,
    start_ns: int,
    end_ns: int,
    value,
) -> float:
    if end_ns <= start_ns:
        raise ValueError("time-weighted interval must be positive")
    selected = [
        record
        for record in records
        if start_ns <= int(record["submit_mono_ns"]) <= end_ns
    ]
    if not selected:
        raise ValueError("time-weighted interval contains no records")
    selected.sort(key=lambda record: int(record["submit_mono_ns"]))
    cursor = start_ns
    total = 0.0
    current = _as_number(value(selected[0]), name="time-weighted value")
    for record in selected[1:]:
        timestamp = int(record["submit_mono_ns"])
        if timestamp < cursor:
            raise ValueError("telemetry timestamps are not monotonic")
        total += current * (timestamp - cursor)
        cursor = timestamp
        current = _as_number(value(record), name="time-weighted value")
    total += current * (end_ns - cursor)
    return total / (end_ns - start_ns)


def summarize_engine(
    records: Sequence[Mapping[str, Any]],
    *,
    start_ns: int,
    end_ns: int,
    request_count: int,
) -> dict[str, Any]:
    """Summarize a measured engine interval from Layer-1 records."""

    if request_count <= 0:
        raise ValueError("request_count must be positive")
    layer1 = [record for record in records if "step_index" in record]
    if not layer1:
        raise ValueError("engine stream has no Layer-1 records")
    step_indexes = [int(record["step_index"]) for record in layer1]
    if len(step_indexes) != len(set(step_indexes)):
        raise ValueError("engine Layer-1 step indexes are not unique")
    if step_indexes != sorted(step_indexes):
        raise ValueError("engine Layer-1 step indexes are not ordered")
    if any(int(record.get("dropped_records_before", 0)) != 0 for record in layer1):
        raise ValueError("engine Layer-1 stream reports dropped records")

    interval = [
        record
        for record in layer1
        if start_ns <= int(record["submit_mono_ns"]) <= end_ns
    ]
    if not interval:
        raise ValueError("engine interval has no Layer-1 records")
    executed = [record for record in interval if bool(record["model_executed"])]
    if not executed:
        raise ValueError("engine interval has no executed model steps")
    duration_s = (end_ns - start_ns) / 1e9
    if duration_s <= 0:
        raise ValueError("engine interval duration must be positive")

    batch_sizes = [int(record["scheduled_requests"]) for record in executed]
    prefill_tokens = [int(record["prefill_tokens"]) for record in executed]
    decode_tokens = [int(record["decode_tokens"]) for record in executed]
    batch_tokens = [
        prefill + decode
        for prefill, decode in zip(prefill_tokens, decode_tokens, strict=True)
    ]
    decode_batches = [int(record["decode_batch_size"]) for record in executed]
    if any(value < 0 for value in batch_sizes + batch_tokens + decode_batches):
        raise ValueError("engine batch counters must be non-negative")
    if any(
        int(record["prefill_tokens"]) + int(record["decode_tokens"])
        <= 0
        for record in executed
    ):
        raise ValueError("executed engine step has no scheduled tokens")

    waiting_mean = _time_weighted_mean(
        interval,
        start_ns=start_ns,
        end_ns=end_ns,
        value=lambda record: record["queues"]["waiting"],
    )
    running_mean = _time_weighted_mean(
        interval,
        start_ns=start_ns,
        end_ns=end_ns,
        value=lambda record: record["queues"]["running"],
    )
    kv_mean = _time_weighted_mean(
        interval,
        start_ns=start_ns,
        end_ns=end_ns,
        value=lambda record: record["kv"]["usage"],
    )
    kv_values = [float(record["kv"]["usage"]) for record in interval]
    if any(not 0.0 <= value <= 1.0 for value in kv_values):
        raise ValueError("engine KV usage must be in [0, 1]")
    total_prefill = sum(prefill_tokens)
    total_decode = sum(decode_tokens)
    graph_modes = [str(record["cudagraph"]["runtime_mode"]) for record in executed]
    bucket_tokens = sum(int(record["cudagraph"]["bucket_tokens"]) for record in executed)
    padding_tokens = sum(int(record["cudagraph"]["padding_tokens"]) for record in executed)

    common = {
        "scheduler_steps_per_s": len(executed) / duration_s,
        "batch_size": numeric(batch_sizes),
        "batch_tokens": numeric(batch_tokens),
        "decode_batch_size": numeric(decode_batches),
        "prefill_token_fraction": total_prefill / (total_prefill + total_decode),
        "queue_waiting_mean": waiting_mean,
        "queue_running_mean": running_mean,
        "queue_waiting_time_per_request_ms": waiting_mean * duration_s * 1000.0 / request_count,
        "queue_running_time_per_request_ms": running_mean * duration_s * 1000.0 / request_count,
        "preemptions": sum(int(record["preemptions"]) for record in executed),
    }
    result = {
        "schema": SCHEMA,
        "source": "engine_layer1",
        "interval": {
            "start_ns": start_ns,
            "end_ns": end_ns,
            "duration_s": duration_s,
            "request_count": request_count,
        },
        "common": common,
        "engine_only": {
            "kv_usage_mean": kv_mean,
            "kv_usage_max": max(kv_values),
            "kv_usage_end_minus_start": kv_values[-1] - kv_values[0],
            "graph_none_share": graph_modes.count("NONE") / len(graph_modes),
            "graph_full_share": graph_modes.count("FULL") / len(graph_modes),
            "graph_padding_fraction": padding_tokens / max(1, bucket_tokens),
        },
        "simulator_only": {},
        "sanity": {
            "records": len(interval),
            "executed_steps": len(executed),
            "step_index_min": min(int(record["step_index"]) for record in interval),
            "step_index_max": max(int(record["step_index"]) for record in interval),
            "invariants": {
                "positive_duration": duration_s > 0,
                "positive_request_count": request_count > 0,
                "zero_drops": True,
                "nonnegative_counters": True,
                "kv_bounded": True,
                "batch_values_not_all_identical": any(
                    summary["distinct_n"] > 1
                    for summary in (
                        common["batch_size"],
                        common["batch_tokens"],
                        common["decode_batch_size"],
                    )
                ),
            },
        },
    }
    return result


def _csv_rows(path: Path) -> list[dict[str, str]]:
    with path.open(encoding="utf-8", newline="") as source:
        rows = list(csv.DictReader(source))
    if not rows:
        raise ValueError(f"{path}: CSV contains no rows")
    return rows


def _column(rows: Sequence[Mapping[str, str]], name: str) -> list[float]:
    if name not in rows[0]:
        raise ValueError(f"CSV is missing required column {name!r}")
    values = []
    for row in rows:
        text = row.get(name, "")
        if text == "":
            raise ValueError(f"CSV column {name!r} contains an empty value")
        values.append(_as_number(float(text), name=name))
    return values


def summarize_frontier(
    *,
    system_metrics_path: Path,
    request_metrics_path: Path,
    batch_metrics_path: Path | None = None,
    ledger_path: Path | None = None,
) -> dict[str, Any]:
    """Summarize a Frontier run, retaining unavailable state as null."""

    system = json.loads(system_metrics_path.read_text(encoding="utf-8"))
    throughput = system["throughput_metrics"]
    duration_s = _as_number(
        throughput["total_duration_seconds"], name="total_duration_seconds"
    )
    if duration_s <= 0:
        raise ValueError("Frontier duration must be positive")
    request_rows = _csv_rows(request_metrics_path)
    waiting_ms = _column(request_rows, "request_waiting_time_total")
    e2e_ms = _column(request_rows, "request_e2e_time")
    running_ms = [max(0.0, e2e - waiting) for e2e, waiting in zip(e2e_ms, waiting_ms, strict=True)]
    duration_ms = duration_s * 1000.0
    request_count = len(request_rows)

    common: dict[str, Any] = {
        "scheduler_steps_per_s": None,
        "batch_size": None,
        "batch_tokens": None,
        "decode_batch_size": None,
        "prefill_token_fraction": None,
        "queue_waiting_mean": sum(waiting_ms) / duration_ms,
        "queue_running_mean": sum(running_ms) / duration_ms,
        "queue_waiting_time_per_request_ms": fmean(waiting_ms),
        "queue_running_time_per_request_ms": fmean(running_ms),
        "preemptions": sum(
            int(float(row.get("request_total_preemption_count") or 0))
            for row in request_rows
        ),
    }
    batch_rows: list[dict[str, str]] = []
    if batch_metrics_path is not None:
        batch_rows = _csv_rows(batch_metrics_path)
        batch_sizes = _column(batch_rows, "batch_size")
        batch_tokens = _column(batch_rows, "batch_num_tokens")
        prefill_tokens = _column(batch_rows, "batch_num_prefill_tokens")
        decode_tokens = _column(batch_rows, "batch_num_decode_tokens")
        if any(value < 0 for value in batch_sizes + batch_tokens + prefill_tokens + decode_tokens):
            raise ValueError("Frontier batch counters must be non-negative")
        common.update(
            {
                "scheduler_steps_per_s": len(batch_rows) / duration_s,
                "batch_size": numeric(batch_sizes),
                "batch_tokens": numeric(batch_tokens),
                "decode_batch_size": numeric(decode_tokens),
                "prefill_token_fraction": sum(prefill_tokens)
                / max(1.0, sum(prefill_tokens) + sum(decode_tokens)),
            }
        )

    ledger_rows: list[dict[str, Any]] = []
    if ledger_path is not None:
        ledger_rows = load_jsonl(ledger_path)
        for row in ledger_rows:
            start = _as_number(row["stage_start_ts"], name="stage_start_ts")
            end = _as_number(row["stage_end_ts"], name="stage_end_ts")
            if end < start:
                raise ValueError("Frontier ledger has a negative stage duration")

    batch_distinct = (
        max(
            summary["distinct_n"]
            for summary in (
                common["batch_size"],
                common["batch_tokens"],
                common["decode_batch_size"],
            )
        )
        if batch_rows
        else None
    )
    return {
        "schema": SCHEMA,
        "source": "frontier",
        "interval": {
            "duration_s": duration_s,
            "request_count": request_count,
        },
        "common": common,
        "engine_only": {
            "kv_usage_mean": None,
            "kv_usage_max": None,
            "kv_usage_end_minus_start": None,
            "graph_none_share": None,
            "graph_full_share": None,
            "graph_padding_fraction": None,
        },
        "simulator_only": {
            "request_waiting_time_ms": numeric(waiting_ms),
            "request_running_time_ms": numeric(running_ms),
            "ledger_rows": len(ledger_rows) if ledger_path is not None else None,
        },
        "sanity": {
            "request_rows": request_count,
            "batch_rows": len(batch_rows),
            "ledger_rows": len(ledger_rows),
            "invariants": {
                "positive_duration": duration_s > 0,
                "positive_request_count": request_count > 0,
                "nonnegative_counters": True,
                "request_values_not_all_identical": max(
                    numeric(waiting_ms)["distinct_n"],
                    numeric(running_ms)["distinct_n"],
                )
                > 1,
                "batch_values_not_all_identical": (
                    batch_distinct > 1 if batch_distinct is not None else None
                ),
            },
        },
    }


def residual(real: Mapping[str, Any], simulated: Mapping[str, Any]) -> dict[str, Any]:
    if real.get("schema") != SCHEMA or simulated.get("schema") != SCHEMA:
        raise ValueError("residual inputs must use the common-state schema")
    values = {}
    missing = []
    for name, real_value in real["common"].items():
        sim_value = simulated["common"].get(name)
        if isinstance(real_value, dict):
            if not isinstance(sim_value, dict):
                missing.append(name)
                continue
            for statistic in ("mean", "max", "cv"):
                key = f"{name}.{statistic}"
                values[key] = float(real_value[statistic]) - float(sim_value[statistic])
            continue
        if real_value is None or sim_value is None:
            missing.append(name)
            continue
        values[name] = float(real_value) - float(sim_value)
    return {
        "schema": "telemetry-state-residual-v1",
        "values": values,
        "missing_common_fields": sorted(missing),
        "coverage": {
            "available": len(values),
            "missing": len(missing),
            "common_field_count": len(real["common"]),
        },
    }