#!/usr/bin/env python3

from __future__ import annotations

import asyncio
import importlib.util
import json
import sys
import tempfile
import unittest
from unittest import mock
from pathlib import Path
from types import SimpleNamespace

from aiohttp import web


HERE = Path(__file__).parent


def load_module(name: str, filename: str):
    spec = importlib.util.spec_from_file_location(name, HERE / filename)
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


client = load_module("phase3_client", "opprof_phase3_client.py")
controller = load_module("phase3_controller", "opprof_phase3_controller.py")
sys.modules["opprof_phase3_controller"] = controller
matrix = load_module("phase3_matrix", "opprof_phase3_matrix.py")


def rows(n: int, arrival: str = "steady") -> list[dict]:
    return [
        {
            "schema": 1,
            "request_id": f"T-{i:05d}",
            "pattern_id": "T",
            "kind": "synthetic",
            "input_tokens": 8,
            "output_tokens": 2,
            "arrival": arrival,
            "token_seed": i + 1,
        }
        for i in range(n)
    ]


class MockServer:
    def __init__(self, delay: float = 0.01, trace_dir: Path | None = None) -> None:
        self.active = 0
        self.max_active = 0
        self.payloads = []
        self.runner = None
        self.port = None
        self.delay = delay
        self.trace_dir = trace_dir
        self.profile_count = 0

    async def completion(self, request):
        payload = await request.json()
        self.payloads.append(payload)
        self.active += 1
        self.max_active = max(self.max_active, self.active)
        await asyncio.sleep(self.delay)
        response = web.StreamResponse(
            status=200, headers={"Content-Type": "text/event-stream"}
        )
        await response.prepare(request)
        await response.write(
            b'data: {"choices":[{"text":"x"}],"usage":null}\n\n'
        )
        usage = json.dumps(
            {
                "choices": [],
                "usage": {
                    "prompt_tokens": len(payload["prompt"]),
                    "completion_tokens": payload["max_tokens"],
                },
            }
        ).encode()
        await response.write(b"data: " + usage + b"\n\n")
        await response.write(b"data: [DONE]\n\n")
        await response.write_eof()
        self.active -= 1
        return response

    async def start_profile(self, request):
        self.profile_count += 1
        path = self.trace_dir / f"window-{self.profile_count}.pt.trace.json"
        path.write_text('{"traceEvents": []}')
        return web.Response(status=200)

    async def stop_profile(self, request):
        return web.Response(status=200)

    async def start(self):
        app = web.Application()
        app.router.add_post("/v1/completions", self.completion)
        if self.trace_dir is not None:
            app.router.add_post("/start_profile", self.start_profile)
            app.router.add_post("/stop_profile", self.stop_profile)
        self.runner = web.AppRunner(app)
        await self.runner.setup()
        site = web.TCPSite(self.runner, "127.0.0.1", 0)
        await site.start()
        self.port = site._server.sockets[0].getsockname()[1]

    async def stop(self):
        await self.runner.cleanup()


def run_args(
    manifest: Path,
    result_dir: Path,
    port: int,
    load_point: str = "saturation",
    saturation_result: Path | None = None,
):
    return SimpleNamespace(
        manifest=str(manifest),
        base_url=f"http://127.0.0.1:{port}",
        model="mock",
        load_point=load_point,
        request_rate="inf" if load_point == "saturation" else None,
        saturation_result=None if saturation_result is None else str(saturation_result),
        rate_fraction=0.60,
        max_concurrency=3,
        ignore_eos=True,
        temperature=0,
        warmup_seconds=0.04,
        clean_segment_seconds=0.04,
        num_clean_segments=3,
        profile_after_clean=False,
        num_profile_windows=0,
        profile_warmup_iterations=2,
        profile_active_iterations=8,
        profile_trace_dir=None,
        profile_timeout_seconds=1,
        recovery_seconds=0,
        post_clean_seconds=0,
        drain_timeout_seconds=1,
        workload_seed=20260712,
        server_seed=20260712,
        result_dir=str(result_dir),
    )


class Phase3ToolTests(unittest.TestCase):
    def write_warmup_stream(self, root: Path, tokens: list[int]) -> None:
        stream_dir = root / "opprof"
        stream_dir.mkdir()
        records = []
        step = 0
        for bin_index, token_count in enumerate(tokens):
            per_step, remainder = divmod(token_count, 16)
            for in_bin in range(16):
                records.append(
                    {
                        "step_index": step,
                        "model_executed": True,
                        "submit_mono_ns": int(
                            1e9
                            + (45 + bin_index * 5 + (in_bin + 0.5) / 16 * 5)
                            * 1e9
                        ),
                        "prefill_tokens": per_step + (in_bin < remainder),
                        "decode_tokens": 0,
                    }
                )
                step += 1
        (stream_dir / "test.jsonl").write_text(
            "".join(json.dumps(record) + "\n" for record in records)
        )

    def test_p10_warmup_stability_accepts_flat_trailing_quartile(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            self.write_warmup_stream(root, [1600, 1600, 1600])
            result = matrix.p10_warmup_stability(root, int(1e9))
            self.assertTrue(result["passed"])
            self.assertEqual(result["step_counts"], [16, 16, 16])
            self.assertEqual(result["scheduled_token_throughput"], [320, 320, 320])
            self.assertEqual(result["normalized_drift"], 0)

    def test_p10_warmup_stability_rejects_large_drift(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            self.write_warmup_stream(root, [1600, 3200, 4800])
            result = matrix.p10_warmup_stability(root, int(1e9))
            self.assertFalse(result["passed"])
            self.assertGreater(result["normalized_drift"], 0.10)

    def test_only_clean_window_request_failures_are_hard_failures(self):
        requests = [
            {"success": False, "completed_s": 59.9, "error_kind": "warmup"},
            {"success": False, "completed_s": 60.0, "error_kind": "clean"},
            {"success": False, "completed_s": 299.9, "error_kind": "clean"},
            {"success": False, "completed_s": 300.0, "error_kind": "recovery"},
            {"success": True, "completed_s": 100.0, "error_kind": None},
        ]
        summary = matrix.summarize_request_failures(requests, 60.0, 300.0)
        self.assertEqual(summary["failed"], 4)
        self.assertEqual(summary["clean_failed"], 2)
        self.assertEqual(summary["excluded"], 2)
        self.assertEqual(summary["excluded_kinds"], {"warmup": 1, "recovery": 1})

    def test_synthetic_prompt_exact_and_unique_early_prefix(self):
        generated = [client.synthetic_prompt(row) for row in rows(200)]
        self.assertTrue(all(len(value) == 8 for value in generated))
        self.assertEqual(len({tuple(value[:3]) for value in generated}), 200)

    def test_p05_manifest_has_exact_half_modes(self):
        with tempfile.TemporaryDirectory() as tmp:
            out = Path(tmp) / "P05.jsonl"
            args = SimpleNamespace(
                id="P05",
                kind="synthetic",
                input_uniform=None,
                input_fixed=None,
                input_mixture='{"uniform:128:512":0.5,"uniform:4096:8192":0.5}',
                output_fixed=64,
                prefix="none",
                arrival="steady",
                num_requests=100,
                workload_seed=20260712,
                num_prefixes=0,
                prefix_len=0,
                suffix_fixed=0,
                out=str(out),
            )
            client.materialize(args)
            manifest = client.load_manifest(out)
            self.assertEqual(sum(r["input_tokens"] <= 512 for r in manifest), 50)
            self.assertEqual(sum(r["input_tokens"] >= 4096 for r in manifest), 50)

    def test_prefix_pool_balanced(self):
        with tempfile.TemporaryDirectory() as tmp:
            out = Path(tmp) / "P08.jsonl"
            args = SimpleNamespace(
                id="P08",
                kind="prefix-pool",
                input_uniform=None,
                input_fixed=None,
                input_mixture=None,
                output_fixed=512,
                prefix="none",
                arrival="burst:8",
                num_requests=80,
                workload_seed=20260712,
                num_prefixes=8,
                prefix_len=1024,
                suffix_fixed=256,
                out=str(out),
            )
            client.materialize(args)
            manifest = client.load_manifest(out)
            counts = {i: 0 for i in range(8)}
            for row in manifest:
                counts[row["prefix_id"]] += 1
                self.assertEqual(row["input_tokens"], 1280)
            self.assertEqual(set(counts.values()), {10})

    def test_fixed_duration_saturation_and_redaction(self):
        async def case():
            with tempfile.TemporaryDirectory() as tmp:
                root = Path(tmp)
                manifest = root / "m.jsonl"
                client.atomic_jsonl(manifest, rows(1000))
                mock = MockServer()
                await mock.start()
                try:
                    result = await client.run_load(
                        run_args(manifest, root / "result", mock.port)
                    )
                finally:
                    await mock.stop()
                self.assertEqual(result["max_in_flight"], 3)
                self.assertEqual(mock.max_active, 3)
                self.assertAlmostEqual(result["clean"]["duration_s"], 0.12, places=9)
                self.assertEqual(len(result["segments"]), 3)
                self.assertLessEqual(result["drain_seconds"], 1)
                self.assertTrue(
                    all(p["max_tokens"] == 2 and p["ignore_eos"] for p in mock.payloads)
                )
                text = (root / "result/requests.jsonl").read_text()
                self.assertNotIn('"prompt":', text)
                self.assertNotIn('"text":', text)

        asyncio.run(case())

    def test_profile_control_plane_is_not_starved_by_data_connector(self):
        async def case():
            with tempfile.TemporaryDirectory() as tmp:
                root = Path(tmp)
                manifest = root / "m.jsonl"
                trace_dir = root / "traces"
                trace_dir.mkdir()
                client.atomic_jsonl(manifest, rows(1000))
                server = MockServer(delay=0.5, trace_dir=trace_dir)
                await server.start()
                args = run_args(manifest, root / "result", server.port)
                args.max_concurrency = 1
                args.warmup_seconds = 0.01
                args.clean_segment_seconds = 0.01
                args.num_clean_segments = 1
                args.profile_after_clean = True
                args.num_profile_windows = 1
                args.profile_trace_dir = str(trace_dir)
                args.recovery_seconds = 0
                try:
                    result = await client.run_load(args)
                finally:
                    await server.stop()
                self.assertEqual(server.max_active, 1)
                self.assertEqual(server.profile_count, 1)
                self.assertLess(result["profiles"][0]["start_return_s"], 0.25)

        asyncio.run(case())

    def test_drain_timeout_is_a_hard_sanity_failure(self):
        async def case():
            with tempfile.TemporaryDirectory() as tmp:
                root = Path(tmp)
                manifest = root / "m.jsonl"
                client.atomic_jsonl(manifest, rows(1000))
                mock = MockServer()
                await mock.start()
                try:
                    args = run_args(manifest, root / "result", mock.port)
                    args.drain_timeout_seconds = 0
                    with self.assertRaisesRegex(RuntimeError, "drain_within_timeout"):
                        await client.run_load(args)
                finally:
                    await mock.stop()
                sanity = json.loads((root / "result/sanity.json").read_text())
                self.assertFalse(sanity["invariants"]["drain_within_timeout"])

        asyncio.run(case())

    def test_burst_schedule_is_eight_at_once(self):
        async def case():
            with tempfile.TemporaryDirectory() as tmp:
                root = Path(tmp)
                manifest = root / "m.jsonl"
                client.atomic_jsonl(manifest, rows(1000, "burst:8"))
                sat = root / "sat.json"
                client.atomic_json(
                    sat, {"clean": {"completed_throughput_rps": 100.0}}
                )
                mock = MockServer()
                await mock.start()
                try:
                    args = run_args(
                        manifest, root / "result", mock.port, "moderate", sat
                    )
                    args.warmup_seconds = 0
                    args.clean_segment_seconds = 0.4 / 3
                    result = await client.run_load(args)
                finally:
                    await mock.stop()
                records = [
                    json.loads(line)
                    for line in (root / "result/requests.jsonl").read_text().splitlines()
                ]
                groups = {}
                for record in records:
                    groups.setdefault(round(record["scheduled_s"], 6), 0)
                    groups[round(record["scheduled_s"], 6)] += 1
                self.assertTrue(all(value == 8 for value in groups.values()))
                starts = sorted(groups)
                if len(starts) > 1:
                    self.assertAlmostEqual(starts[1] - starts[0], 8 / 60, places=5)
                self.assertAlmostEqual(result["request_rate"], 60.0)

        asyncio.run(case())

    def test_manifest_exhaustion_stops_instead_of_wrapping(self):
        async def case():
            with tempfile.TemporaryDirectory() as tmp:
                root = Path(tmp)
                manifest = root / "m.jsonl"
                client.atomic_jsonl(manifest, rows(2))
                mock = MockServer()
                await mock.start()
                try:
                    with self.assertRaises(client.ManifestExhausted):
                        await client.run_load(
                            run_args(manifest, root / "result", mock.port)
                        )
                finally:
                    await mock.stop()

        asyncio.run(case())

    def test_cpu_affinity_sets_are_disjoint_and_cover_host(self):
        values = []
        for spec in controller.CPU_MAP.values():
            lo, hi = [int(x) for x in spec.split("-")]
            values.extend(range(lo, hi + 1))
        self.assertEqual(sorted(values), list(range(160)))
        self.assertEqual(len(values), len(set(values)))

    def test_preflight_waits_for_three_stable_zero_samples(self):
        def sample(memory):
            return [
                {
                    "index": 0,
                    "uuid": "GPU-test",
                    "memory_used_mib": memory,
                    "utilization_pct": 0,
                }
            ]

        with tempfile.TemporaryDirectory() as tmp:
            with (
                mock.patch.object(
                    controller,
                    "gpu_query",
                    side_effect=[sample(4), sample(0), sample(0), sample(0)],
                ),
                mock.patch.object(controller, "compute_apps", return_value=[]),
                mock.patch.object(controller, "run_text", return_value=""),
                mock.patch.object(controller.time, "sleep"),
            ):
                root = Path(tmp)
                controller.preflight([0], root)
                samples = json.loads((root / "gpu-before-samples.json").read_text())
                self.assertEqual(len(samples), 4)
                self.assertEqual(samples[0]["gpus"][0]["memory_used_mib"], 4)
                self.assertEqual(samples[-1]["gpus"][0]["memory_used_mib"], 0)

    def test_kernel_mapping_priority(self):
        cases = {
            "void vllm::moe::topkGating": "moe_router",
            "ncclDevKernel_AllReduce": "collective",
            "flash_fwd_kernel": "attention",
            "nvjet_sm90_tst": "moe_gemm",
            "argmax_kernel": "sampler",
            "cutlass_gemm": "dense_gemm",
            "triton_red_fused_add_rms_norm": "norm_elementwise",
            "cache_swap_kernel": "kv_memory",
            "unknown": "other",
        }
        self.assertEqual(
            {name: controller.classify_kernel(name) for name in cases}, cases
        )

    def test_atomic_state_replacement(self):
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "state.json"
            controller.atomic_json(path, {"schema": 1, "value": 1})
            controller.atomic_json(path, {"schema": 1, "value": 2})
            self.assertEqual(json.loads(path.read_text())["value"], 2)
            self.assertFalse(list(path.parent.glob("*.tmp.*")))

    def test_fingerprint_survives_json_roundtrip(self):
        original_run_text = controller.run_text
        original_hash = controller.sha256_file
        try:
            controller.run_text = lambda *args, **kwargs: "deadbeef\n"
            controller.sha256_file = lambda path: "a" * 64
            fingerprint = controller.make_fingerprint()
        finally:
            controller.run_text = original_run_text
            controller.sha256_file = original_hash
        self.assertEqual(json.loads(json.dumps(fingerprint)), fingerprint)
        self.assertEqual(set(fingerprint["cpu_map"]), {str(i) for i in range(8)})

    def test_server_shutdown_signals_parent_before_group(self):
        process = SimpleNamespace(pid=12345, poll=lambda: None)
        with (
            mock.patch.object(controller.os, "kill") as kill,
            mock.patch.object(controller.os, "killpg") as killpg,
            mock.patch.object(controller, "_process_group_alive", return_value=False),
        ):
            controller.stop_servers([process])
        kill.assert_called_once_with(12345, controller.signal.SIGINT)
        killpg.assert_not_called()


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