#!/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")


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) -> None:
        self.active = 0
        self.max_active = 0
        self.payloads = []
        self.runner = None
        self.port = None

    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(0.01)
        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(self):
        app = web.Application()
        app.router.add_post("/v1/completions", self.completion)
        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 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_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_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)