#!/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__).parent


def load(name: str):
    path = ROOT / name
    spec = importlib.util.spec_from_file_location(path.stem, path)
    assert spec and spec.loader
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


class FidelityEnvelopeTest(unittest.TestCase):
    def test_block_identities_are_parent_sensitive_and_prefix_stable(self) -> None:
        module = load("prepare_exact_trace.py")
        prefix = list(range(32))
        left = module.block_identities(prefix + [100, 101], 16)
        right = module.block_identities(prefix + [200, 201], 16)
        self.assertEqual(left[:2], right[:2])
        self.assertNotEqual(left[2], right[2])
        changed_parent = module.block_identities([999] + prefix[1:] + [100, 101], 16)
        self.assertNotEqual(left[0], changed_parent[0])
        self.assertNotEqual(left[1], changed_parent[1])

    def test_root_sessions_follow_parent_chain(self) -> None:
        module = load("prepare_exact_trace.py")
        rows = [
            {"chat_id": 10, "parent_chat_id": -1},
            {"chat_id": 11, "parent_chat_id": 10},
            {"chat_id": 12, "parent_chat_id": 11},
            {"chat_id": 20, "parent_chat_id": -1},
        ]
        self.assertEqual(module.root_sessions(rows), {10: 10, 11: 10, 12: 10, 20: 20})

    def test_materialize_allreduce(self) -> None:
        module = load("materialize_frontier_allreduce.py")
        rows = []
        for tp in (2, 4):
            for tokens in (1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192):
                rows.append(
                    {
                        "tensor_parallel_size": tp,
                        "num_tokens": tokens,
                        "hidden_dim": 2048,
                        "payload_bytes": tokens * 2048 * 2,
                        "critical_path_median_ms": tp + tokens / 1000,
                    }
                )
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            source = root / "allreduce.json"
            source.write_text(
                json.dumps(
                    {
                        "schema_version": "qwen30_vllm020_allreduce_frozen.v1",
                        "rows": rows,
                    }
                )
            )
            output = root / "all_reduce.csv"
            manifest = module.convert(source, output)
            self.assertEqual(manifest["rows"], 24)
            with output.open(newline="") as handle:
                converted = list(csv.DictReader(handle))
            self.assertEqual(converted[0]["num_workers"], "2")
            self.assertEqual(converted[0]["size"], "4096")
            self.assertEqual(converted[-1]["num_workers"], "4")
            self.assertEqual(converted[-1]["size"], str(8192 * 2048 * 2))
            self.assertEqual(
                converted[-1]["time_stats.all_reduce.median"],
                str(4 + 8192 / 1000),
            )


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