#!/usr/bin/env python3
from __future__ import annotations

import csv
import json
import math
import tempfile
from pathlib import Path

from common_state import numeric, residual, summarize_engine, summarize_frontier
from run_frontier_state import enable_state_outputs


def engine_record(
    step: int,
    timestamp_ns: int,
    *,
    batch: int,
    prefill: int,
    decode: int,
    waiting: int,
    running: int,
    kv: float,
) -> dict[str, object]:
    return {
        "step_index": step,
        "submit_mono_ns": timestamp_ns,
        "model_executed": True,
        "scheduled_requests": batch,
        "decode_batch_size": decode,
        "prefill_tokens": prefill,
        "decode_tokens": decode,
        "preemptions": 0,
        "queues": {"waiting": waiting, "running": running},
        "kv": {"usage": kv},
        "cudagraph": {
            "runtime_mode": "NONE" if step == 0 else "FULL",
            "bucket_tokens": prefill + decode + 1,
            "padding_tokens": 1,
        },
        "dropped_records_before": 0,
    }


def main() -> None:
    assert numeric((-2.0, -1.0))["cv"] >= 0.0
    command = [
        "python3",
        "-m",
        "frontier.main",
        "--no-metrics_config_store_frontier_stage_batch_ledger",
    ]
    state_command = enable_state_outputs(command)
    assert "--metrics_config_store_frontier_stage_batch_ledger" in state_command
    assert "--metrics_config_keep_individual_batch_metrics" in state_command
    assert "--no-metrics_config_store_frontier_stage_batch_ledger" not in state_command

    engine = summarize_engine(
        [
            engine_record(
                0,
                0,
                batch=2,
                prefill=6,
                decode=0,
                waiting=1,
                running=2,
                kv=0.1,
            ),
            engine_record(
                1,
                1_000_000_000,
                batch=4,
                prefill=0,
                decode=4,
                waiting=3,
                running=4,
                kv=0.3,
            ),
        ],
        start_ns=0,
        end_ns=2_000_000_000,
        request_count=2,
    )
    assert math.isclose(engine["common"]["queue_waiting_mean"], 2.0)
    assert math.isclose(engine["common"]["queue_running_mean"], 3.0)
    assert math.isclose(
        engine["common"]["queue_waiting_time_per_request_ms"], 2000.0
    )
    assert engine["common"]["batch_size"]["distinct_n"] == 2
    assert math.isclose(engine["common"]["prefill_token_fraction"], 0.6)

    with tempfile.TemporaryDirectory() as temporary:
        root = Path(temporary)
        system = root / "system.json"
        requests = root / "requests.csv"
        batches = root / "batches.csv"
        ledger = root / "ledger.jsonl"
        system.write_text(
            json.dumps(
                {
                    "throughput_metrics": {
                        "total_duration_seconds": 2.0,
                    }
                }
            ),
            encoding="utf-8",
        )
        with requests.open("w", encoding="utf-8", newline="") as output:
            writer = csv.DictWriter(
                output,
                fieldnames=[
                    "request_e2e_time",
                    "request_waiting_time_total",
                    "request_total_preemption_count",
                ],
            )
            writer.writeheader()
            writer.writerow(
                {
                    "request_e2e_time": 2000,
                    "request_waiting_time_total": 1000,
                    "request_total_preemption_count": 0,
                }
            )
            writer.writerow(
                {
                    "request_e2e_time": 3000,
                    "request_waiting_time_total": 3000,
                    "request_total_preemption_count": 0,
                }
            )
        with batches.open("w", encoding="utf-8", newline="") as output:
            writer = csv.DictWriter(
                output,
                fieldnames=[
                    "batch_size",
                    "batch_num_tokens",
                    "batch_num_prefill_tokens",
                    "batch_num_decode_tokens",
                ],
            )
            writer.writeheader()
            writer.writerow(
                {
                    "batch_size": 2,
                    "batch_num_tokens": 6,
                    "batch_num_prefill_tokens": 6,
                    "batch_num_decode_tokens": 0,
                }
            )
            writer.writerow(
                {
                    "batch_size": 4,
                    "batch_num_tokens": 4,
                    "batch_num_prefill_tokens": 0,
                    "batch_num_decode_tokens": 4,
                }
            )
        ledger.write_text(
            json.dumps({"stage_start_ts": 0.0, "stage_end_ts": 1.0}) + "\n"
            + json.dumps({"stage_start_ts": 1.0, "stage_end_ts": 2.0})
            + "\n",
            encoding="utf-8",
        )
        simulator = summarize_frontier(
            system_metrics_path=system,
            request_metrics_path=requests,
            batch_metrics_path=batches,
            ledger_path=ledger,
        )
    assert math.isclose(simulator["common"]["queue_waiting_mean"], 2.0)
    assert math.isclose(simulator["common"]["queue_running_mean"], 0.5)
    assert simulator["common"]["batch_size"]["distinct_n"] == 2
    difference = residual(engine, simulator)
    assert difference["coverage"]["missing"] == 0
    assert difference["coverage"]["available"] == 16
    assert math.isclose(difference["values"]["queue_waiting_mean"], 0.0)
    print("telemetry common state: PASS")


if __name__ == "__main__":
    main()