#!/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
REPO_ROOT = ROOT.parents[1]


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):
    @staticmethod
    def write_allreduce_csv(path: Path, tps: tuple[int, ...]) -> None:
        fields = [
            "time_stats.all_reduce.median",
            "num_workers",
            "size",
            "collective",
        ]
        with path.open("w", newline="") as handle:
            writer = csv.DictWriter(handle, fieldnames=fields, lineterminator="\n")
            writer.writeheader()
            for tp in tps:
                for size in (8192, 65536):
                    writer.writerow(
                        {
                            "time_stats.all_reduce.median": 0.1 * tp,
                            "num_workers": tp,
                            "size": size,
                            "collective": "all_reduce",
                        }
                    )

    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),
            )

    def test_allreduce_profile_requires_every_selected_tp(self) -> None:
        module = load("run_frontier_qwen30_exact_trace_surface.py")
        with tempfile.TemporaryDirectory() as temporary:
            profile = Path(temporary) / "allreduce.csv"
            self.write_allreduce_csv(profile, (4,))
            with self.assertRaisesRegex(
                ValueError, "coverage missing TPs \\[8\\].*fallback is forbidden"
            ):
                module.validate_allreduce_profile(profile, {4, 8})

            self.write_allreduce_csv(profile, (4, 8))
            validated = module.validate_allreduce_profile(profile, {4, 8})
            self.assertEqual(validated["required_tp_coverage"], [4, 8])
            self.assertEqual(validated["rows_by_tp"], {"4": 2, "8": 2})

    def test_collective_analytical_fallback_is_runtime_failure(self) -> None:
        module = load("run_frontier_qwen30_exact_trace_surface.py")
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            (root / "stdout.log").write_text(
                "All-reduce profiling data is empty after filtering\n"
            )
            (root / "stderr.log").write_text("unrelated warning\n")
            evidence = module.collective_fallback_evidence(root)
            self.assertEqual(len(evidence), 1)
            self.assertEqual(evidence[0]["log"], "stdout.log")

    def test_qwen235_state_replay_changes_observation_flags_only(self) -> None:
        module = load("run_qwen235_fixed_pd_state_replay.py")
        base = [
            "/usr/bin/python3",
            "-m",
            "frontier.main",
            "--metrics_config_output_dir",
            "/frozen/metrics",
            "--metrics_config_run_id",
            "frozen",
            "--no-metrics_config_store_frontier_stage_batch_ledger",
            "--attn_tensor_parallel_size",
            "8",
            "--moe_expert_parallel_size",
            "8",
            "--communication_collective_profile_path",
            "/profiles/real-allreduce.csv",
        ]
        transformed = module.transform_command(
            base, metrics_root=Path("/new/metrics"), run_id="state-run"
        )
        self.assertEqual(
            module.option_value(transformed, "--attn_tensor_parallel_size"), "8"
        )
        self.assertEqual(
            module.option_value(transformed, "--communication_collective_profile_path"),
            "/profiles/real-allreduce.csv",
        )
        self.assertIn(
            "--metrics_config_store_frontier_stage_batch_ledger", transformed
        )
        self.assertIn("--metrics_config_keep_individual_batch_metrics", transformed)
        self.assertNotIn(
            "--no-metrics_config_store_frontier_stage_batch_ledger", transformed
        )
        traced = module.transform_command(
            base, metrics_root=Path("/new/traced-metrics"), run_id="traced", op_trace=True
        )
        self.assertIn("--metrics_config_enable_op_level_tracing", traced)
        self.assertEqual(
            module.option_value(traced, "--communication_collective_profile_path"),
            "/profiles/real-allreduce.csv",
        )

    def test_qwen235_state_replay_rejects_implicit_ledger_base(self) -> None:
        module = load("run_qwen235_fixed_pd_state_replay.py")
        base = [
            "/usr/bin/python3",
            "-m",
            "frontier.main",
            "--metrics_config_output_dir",
            "/frozen/metrics",
            "--metrics_config_run_id",
            "frozen",
        ]
        with self.assertRaisesRegex(ValueError, "explicitly disable one full Frontier ledger"):
            module.transform_command(
                base, metrics_root=Path("/new/metrics"), run_id="state-run"
            )

    def test_qwen235_real_state_proxy_parser(self) -> None:
        module = load("analyze_qwen235_fixed_pd_state.py")
        line = (
            "Engine 000: Avg prompt throughput: 3276.8 tokens/s, "
            "Avg generation throughput: 204.8 tokens/s, Running: 5 reqs, "
            "Waiting: 0 reqs, GPU KV cache usage: 5.6%, Prefix cache hit rate: 0.0%\n"
        )
        with tempfile.TemporaryDirectory() as temporary:
            path = Path(temporary) / "server.log"
            path.write_text(line)
            self.assertEqual(
                module.parse_real_log(path),
                [
                    {
                        "prompt_tokens_per_s": 3276.8,
                        "generation_tokens_per_s": 204.8,
                        "running": 5,
                        "waiting": 0,
                        "kv_percent": 5.6,
                    }
                ],
            )

    def test_qwen235_iteration_state_parser_excludes_warmup(self) -> None:
        module = load("analyze_qwen235_fixed_pd_state.py")
        rows = [
            "Iteration(0): 1 context requests, 512 context tokens, 0 generation requests, 0 generation tokens, iteration elapsed time: 10.00 ms",
            "Iteration(1): 1 context requests, 4096 context tokens, 4 generation requests, 4 generation tokens, iteration elapsed time: 30.00 ms",
            "Iteration(2): 0 context requests, 0 context tokens, 5 generation requests, 5 generation tokens, iteration elapsed time: 20.00 ms",
        ]
        with tempfile.TemporaryDirectory() as temporary:
            path = Path(temporary) / "server.log"
            path.write_text("\n".join(rows) + "\n")
            parsed = module.parse_iteration_log(path)
        self.assertEqual([row["index"] for row in parsed], [1, 2])
        self.assertEqual(parsed[0]["context_tokens"], 4096)
        self.assertEqual(parsed[1]["generation_requests"], 5)

    def test_qwen235_component_categories_are_additive(self) -> None:
        module = load("analyze_qwen235_fixed_pd_state.py")
        components = {
            name: 1.0
            for names in module.CATEGORIES.values()
            for name in names
        }
        categorized = module.categorized_components(components)
        self.assertEqual(sum(categorized.values()), float(len(components)))
        components["unknown_graph_overhead"] = 1.0
        with self.assertRaisesRegex(ValueError, "component schema drift"):
            module.categorized_components(components)
        for category, names in module.OP_CATEGORIES.items():
            for name in names:
                self.assertEqual(module.op_category(name), category)
        self.assertIsNone(module.op_category("expert_parallel_allreduce_wait"))
        with self.assertRaisesRegex(ValueError, "unclassified"):
            module.op_category("unknown_graph_overhead")
        component_row = {name: 1.0 for name in (*module.CATEGORIES, "total")}
        reweighted = module.reweight(
            {4: {"components_ms": component_row}}, {"4": 3}
        )
        self.assertEqual(reweighted["coverage"], 1.0)
        self.assertEqual(reweighted["components_ms"]["total"], 1.0)

    def test_materialize_qwen235_allreduce_requires_serving_contract(self) -> None:
        module = load("materialize_qwen235_v020_allreduce.py")
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            inputs = []
            for tp in (4, 8):
                path = root / f"tp{tp}.json"
                rows = []
                for tokens in module.TOKEN_POINTS:
                    fusion_limit = {4: 2 * 1024 * 1024, 8: 512 * 1024}[tp]
                    payload_bytes = tokens * module.HIDDEN_DIM * 2
                    median = tp / 10 + tokens / 10000
                    rows.append(
                        {
                            "tensor_parallel_size": tp,
                            "num_tokens": tokens,
                            "hidden_dim": module.HIDDEN_DIM,
                            "payload_bytes": payload_bytes,
                            "selected_backend": (
                                "flashinfer_trtllm_fused_projection"
                                if payload_bytes <= fusion_limit
                                else "pynccl"
                            ),
                            "real_fusion_limit_bytes": fusion_limit,
                            "trials": 3,
                            "repeats_per_trial": 20,
                            "per_trial_rank_samples_ms": [
                                [[median] * 20 for _ in range(tp)] for _ in range(3)
                            ],
                            "critical_path_median_ms": median,
                        }
                    )
                path.write_text(
                    json.dumps(
                        {
                            "schema_version": "vllm020_allreduce_raw.v2",
                            "environment": {
                                "vllm_version": "0.20.0",
                                "vllm_source_commit": module.VLLM_COMMIT,
                                "gpu": "NVIDIA H20",
                                "model": "/models/Qwen3-235B-A22B-FP8",
                                "collective_contract": "qwen235-serving-projected",
                                "disable_custom_all_reduce": True,
                                "backend_env": {
                                    "VLLM_ALLREDUCE_USE_FLASHINFER": "1",
                                    "VLLM_ALLREDUCE_USE_SYMM_MEM": "1",
                                },
                            },
                            "rows": rows,
                        }
                    )
                )
                inputs.append(path)
            output = root / "allreduce.csv"
            manifest_path = root / "manifest.json"
            manifest = module.materialize(inputs, output, manifest_path)
            self.assertEqual(manifest["tp_coverage"], [4, 8])
            self.assertEqual(manifest["rows"], 2 * len(module.TOKEN_POINTS))
            self.assertEqual(
                manifest["observed_backends_by_tp"],
                {
                    "4": ["flashinfer_trtllm_fused_projection", "pynccl"],
                    "8": ["flashinfer_trtllm_fused_projection", "pynccl"],
                },
            )
            with output.open(newline="") as handle:
                rows = list(csv.DictReader(handle))
            self.assertEqual(rows[0]["num_workers"], "4")
            self.assertEqual(rows[-1]["num_workers"], "8")

    def test_batch_profile_retains_one_single_request_anchor_per_tp(self) -> None:
        runner = REPO_ROOT / "runs/frontier-phase-factorial-v0/run_frontier_qwen30_prefill_surface.py"
        spec = importlib.util.spec_from_file_location("qwen30_surface_runner", runner)
        assert spec and spec.loader
        module = importlib.util.module_from_spec(spec)
        sys.modules[spec.name] = module
        spec.loader.exec_module(module)
        profile = ROOT / "profiles/profile-v3-batch-final"
        coverage = module.validate_profile(module.profile_paths(profile))
        self.assertEqual(
            coverage["attention"],
            {
                "1": {"exact_prefill_2048_rows": 1, "profile_batch_size": 1},
                "2": {"exact_prefill_2048_rows": 1, "profile_batch_size": 1},
                "4": {"exact_prefill_2048_rows": 1, "profile_batch_size": 1},
            },
        )

    def test_exact_trace_parser_and_joint_slo_score(self) -> None:
        module = load("run_frontier_qwen30_exact_trace_surface.py")
        self.assertEqual(
            module.classify_frontier_failure(
                "Sequential simulation ended with non-empty scheduler state"
            ),
            "scheduler_stall",
        )
        self.assertEqual(
            module.classify_frontier_failure("unexpected failure"),
            "frontier_error",
        )
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            trace = root / "trace.csv"
            with trace.open("w", newline="") as handle:
                writer = csv.DictWriter(
                    handle,
                    fieldnames=[
                        "arrived_at",
                        "num_prefill_tokens",
                        "num_decode_tokens",
                        "session_id",
                        "block_hash_ids",
                    ],
                    lineterminator="\n",
                )
                writer.writeheader()
                writer.writerows(
                    [
                        {
                            "arrived_at": 1.0,
                            "num_prefill_tokens": 512,
                            "num_decode_tokens": 1,
                            "session_id": 1,
                            "block_hash_ids": "10|11",
                        },
                        {
                            "arrived_at": 2.0,
                            "num_prefill_tokens": 2048,
                            "num_decode_tokens": 2,
                            "session_id": 2,
                            "block_hash_ids": "20|21",
                        },
                    ]
                )
            parsed = module.parse_trace(f"u0p01={trace}", prefix_caching=False)
            self.assertEqual(parsed["requests"], 2)
            self.assertEqual(parsed["shapes"], [(512, 1), (2048, 2)])
            json.dumps(module.trace_manifest_entry(parsed))
            uniform = module.parse_trace(
                f"r1={trace}",
                rate_contract="uniform-spacing",
                prefix_caching=False,
            )
            self.assertEqual(uniform["offered_request_rate"], 1.0)

            metrics = root / "request_metrics.csv"
            with metrics.open("w", newline="") as handle:
                writer = csv.DictWriter(
                    handle,
                    fieldnames=[
                        "Request Id",
                        "request_num_prefill_tokens",
                        "request_num_decode_tokens",
                        "ttft",
                        "request_e2e_time",
                    ],
                    lineterminator="\n",
                )
                writer.writeheader()
                writer.writerows(
                    [
                        {
                            "Request Id": 1,
                            "request_num_prefill_tokens": 2048,
                            "request_num_decode_tokens": 2,
                            "ttft": 1200,
                            "request_e2e_time": 1300,
                        },
                        {
                            "Request Id": 0,
                            "request_num_prefill_tokens": 512,
                            "request_num_decode_tokens": 1,
                            "ttft": 1000,
                            "request_e2e_time": 1000,
                        },
                    ]
                )
            scored = module.score(metrics, parsed["shapes"])
            self.assertEqual(scored["slos"]["tpot_150ms"]["passed"], 2)
            self.assertTrue(scored["slos"]["tpot_150ms"]["feasible"])

    def test_fixed_trace_materializer_uses_disjoint_prefixes(self) -> None:
        module = load("prepare_fixed_frontier_traces.py")
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            manifest = module.materialize(
                root,
                input_tokens=32,
                output_tokens=4,
                requests=2,
                rates=[2.0],
                block_size=16,
            )
            with (root / "r2.csv").open(newline="") as handle:
                rows = list(csv.DictReader(handle))
            self.assertEqual(rows[0]["block_hash_ids"], "1|2")
            self.assertEqual(rows[1]["block_hash_ids"], "3|4")
            self.assertEqual(rows[1]["arrived_at"], "0.500000000000")
            self.assertEqual(manifest["contract"]["prefix_caching"], False)
            steady = module.materialize(
                root / "steady",
                input_tokens=32,
                output_tokens=4,
                requests=2,
                rates=[2.0],
                block_size=16,
                duration_seconds=10.0,
            )
            self.assertEqual(steady["anchors"][0]["requests"], 20)

    def test_exact_real_client_uses_the_prepared_row_vector_contract(self) -> None:
        client = load("qwen30_exact_trace_client.py")
        prepare = load("prepare_exact_trace.py")
        rows = [
            {
                "source_index": 7,
                "arrived_at": 1.25,
                "input_length": 32,
                "output_length": 4,
                "session_id": 3,
                "runtime_block_ids": [11, 12],
            }
        ]
        digest = client.hashlib.sha256()
        prepare.update_digest(digest, [7, 1.25, 32, 4, 3, [11, 12]])
        self.assertEqual(client.row_vector_sha256(rows), digest.hexdigest())
        self.assertEqual(client.ttft_slo_ms(512), 1064.0)


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