#!/usr/bin/env python3

from __future__ import annotations

import csv
import importlib.util
import json
import sys
import tempfile
import unittest
from pathlib import Path


ROOT = Path(__file__).resolve().parent
S3_REAL = ROOT.parent / "frontier-s3-real-v0"


def load_replay_module():
    sys.path.insert(0, str(S3_REAL))
    spec = importlib.util.spec_from_file_location(
        "frontier_prefix_replay", S3_REAL / "run_frontier_prefix_replay.py"
    )
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


class FrontierCodeReplayTest(unittest.TestCase):
    def test_summary_excludes_undefined_single_token_tpot(self) -> None:
        module = load_replay_module()
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            trace = root / "frontier.csv"
            with trace.open("w", newline="") as stream:
                writer = csv.DictWriter(
                    stream,
                    fieldnames=(
                        "arrived_at",
                        "num_prefill_tokens",
                        "num_decode_tokens",
                        "session_id",
                        "block_hash_ids",
                    ),
                )
                writer.writeheader()
                writer.writerow(
                    {
                        "arrived_at": 0,
                        "num_prefill_tokens": 7,
                        "num_decode_tokens": 1,
                        "session_id": 1,
                        "block_hash_ids": "[]",
                    }
                )
                writer.writerow(
                    {
                        "arrived_at": 1,
                        "num_prefill_tokens": 16,
                        "num_decode_tokens": 2,
                        "session_id": 2,
                        "block_hash_ids": 123,
                    }
                )
            metrics = root / "metrics" / "cell"
            metrics.mkdir(parents=True)
            with (metrics / "request_metrics.csv").open("w", newline="") as stream:
                writer = csv.DictWriter(
                    stream,
                    fieldnames=(
                        "request_waiting_time_total",
                        "ttft",
                        "tpot",
                        "request_e2e_time",
                        "request_cached_prefill_tokens",
                        "request_prefix_cache_query_blocks",
                        "request_prefix_cache_hit_blocks",
                    ),
                )
                writer.writeheader()
                writer.writerow(
                    {
                        "request_waiting_time_total": 0,
                        "ttft": 10,
                        "tpot": "",
                        "request_e2e_time": 10,
                        "request_cached_prefill_tokens": 0,
                        "request_prefix_cache_query_blocks": 0,
                        "request_prefix_cache_hit_blocks": 0,
                    }
                )
                writer.writerow(
                    {
                        "request_waiting_time_total": 0,
                        "ttft": 20,
                        "tpot": 5,
                        "request_e2e_time": 25,
                        "request_cached_prefill_tokens": 0,
                        "request_prefix_cache_query_blocks": 1,
                        "request_prefix_cache_hit_blocks": 0,
                    }
                )
            (metrics / "system_metrics.json").write_text("{}")
            (metrics / "frontier_stage_batch_ledger.jsonl").write_text(
                json.dumps({"request_num_tokens": [1]}) + "\n"
            )
            summary = module.summarize(trace, root / "metrics", 2)

        self.assertEqual(summary["requests"], 2)
        self.assertEqual(summary["latency_ms"]["tpot"]["count"], 1)
        self.assertEqual(summary["latency_ms"]["tpot"]["mean"], 5)


if __name__ == "__main__":
    unittest.main()