import json
import os
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path


_ROOT = Path(__file__).resolve().parents[1]
_SRC = _ROOT / "src"

# Allow `python -m unittest` without installing the package.
sys.path.insert(0, str(_SRC))

from llm_mem_planner.cli import (  # noqa: E402
    HardwareSpec,
    Parallelism,
    dtype_bytes_from_str,
    estimate_activation_working_set,
    estimate_kv_cache_gib,
    estimate_params,
    estimate_step_time,
    estimate_weight_params_per_rank,
    load_dims,
)


class TestSmoke(unittest.TestCase):
    def _load_min_dims(self):
        cfg_path = _ROOT / "tests" / "data" / "minimal_dense_config.json"
        cfg = json.loads(cfg_path.read_text())
        return load_dims(cfg)

    def test_parallelism_invalid_raises(self):
        p = Parallelism(tp=8, dp=3, sp=1, ep=1, moe_dp=1)
        with self.assertRaises(ValueError):
            _ = p.attn_tp

    def test_parallelism_tp_zero_raises(self):
        p = Parallelism(tp=0, dp=1, sp=1, ep=1, moe_dp=1)
        with self.assertRaises(ValueError):
            _ = p.attn_tp

    def test_dtype_int4_is_packed(self):
        self.assertEqual(dtype_bytes_from_str("int4"), 0.5)
        self.assertEqual(dtype_bytes_from_str("fp4"), 0.5)

    def test_kv_cache_linear_scaling_seq_len(self):
        d = self._load_min_dims()
        kv_dtype_bytes = 2
        batch = 8.0
        kv1 = estimate_kv_cache_gib(d, batch=batch, seq_len=1024, kv_dtype_bytes=kv_dtype_bytes)
        kv2 = estimate_kv_cache_gib(d, batch=batch, seq_len=2048, kv_dtype_bytes=kv_dtype_bytes)
        self.assertGreater(kv1, 0.0)
        self.assertGreater(kv2, kv1)
        self.assertAlmostEqual(kv2 / kv1, 2.0, delta=0.02)

    def test_memory_functions_run(self):
        d = self._load_min_dims()
        mhc = None
        params = estimate_params(d, mhc=mhc)
        self.assertGreater(params.total, 0)

        parallel = Parallelism(tp=8, dp=2, sp=1, ep=1, moe_dp=1)
        sharded = estimate_weight_params_per_rank(d, params, parallel)
        self.assertGreater(sharded.total, 0.0)

        ws = estimate_activation_working_set(
            d,
            parallel,
            tokens=1024,
            act_dtype_bytes=2,
            mhc_streams=1,
            mhc_branch_streams=1,
        )
        self.assertGreater(ws.peak_gib, 0.0)

    def test_activation_working_set_includes_dense_ffn(self):
        d = self._load_min_dims()
        parallel = Parallelism(tp=1, dp=1, sp=1, ep=1, moe_dp=1)

        ws = estimate_activation_working_set(
            d,
            parallel,
            tokens=1024,
            act_dtype_bytes=2,
            mhc_streams=1,
            mhc_branch_streams=1,
        )
        self.assertGreater(ws.dense_elems, 0)
        self.assertGreater(ws.dense_elems, ws.hidden_elems)
        self.assertEqual(ws.peak_elems, ws.dense_elems)

    def test_perf_step_time_runs(self):
        d = self._load_min_dims()
        parallel = Parallelism(tp=8, dp=2, sp=1, ep=1, moe_dp=1)

        hw = HardwareSpec(
            nodes=1,
            gpus_per_node=8,
            gpu_mem_gib=80.0,
            mem_fraction=1.0,
            misc_gib=0.0,
            gpu_tflops=1000.0,
            gpu_hbm_gibps=2000.0,
            intra_gibps=200.0,
            inter_gibps=50.0,
        )

        step = estimate_step_time(
            d,
            parallel,
            hw=hw,
            phase="decode",
            tokens_per_rank=8,
            seq_len_local=2048,
            global_batch=16,
            weight_dtype_bytes=2,
            act_dtype_bytes=2,
            kv_dtype_bytes=2,
            dp_padding_mode="auto",
        )
        self.assertGreater(step.total_s, 0.0)
        op_names = {op.name for op in step.ops}
        self.assertIn("attn_fused_qkv_a", op_names)
        self.assertIn("kv_cache_write(latent)", op_names)

    def test_cli_help_and_local_config(self):
        cfg_path = _ROOT / "tests" / "data" / "minimal_dense_config.json"
        env = dict(os.environ)
        env["PYTHONPATH"] = str(_SRC)

        help_run = subprocess.run(
            [sys.executable, "-m", "llm_mem_planner", "--help"],
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=False,
        )
        self.assertEqual(help_run.returncode, 0, msg=help_run.stderr)
        self.assertIn("--config", help_run.stdout)

        run = subprocess.run(
            [
                sys.executable,
                "-m",
                "llm_mem_planner",
                "--config",
                str(cfg_path),
                "--seq-len",
                "2048",
                "--batch",
                "16",
                "--tp",
                "8",
                "--dp",
                "2",
                "--sp",
                "1",
                "--ep",
                "1",
            ],
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=False,
        )
        self.assertEqual(run.returncode, 0, msg=run.stderr)
        self.assertIn("== KV Cache", run.stdout)
        self.assertIn("== Parallelism ==", run.stdout)

    def test_cli_unknown_dtype_fails_cleanly(self):
        cfg_path = _ROOT / "tests" / "data" / "minimal_dense_config.json"
        env = dict(os.environ)
        env["PYTHONPATH"] = str(_SRC)

        run = subprocess.run(
            [
                sys.executable,
                "-m",
                "llm_mem_planner",
                "--config",
                str(cfg_path),
                "--seq-len",
                "2048",
                "--batch",
                "16",
                "--tp",
                "8",
                "--dp",
                "2",
                "--sp",
                "1",
                "--ep",
                "1",
                "--weight-dtype",
                "not_a_dtype",
            ],
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=False,
        )
        self.assertNotEqual(run.returncode, 0)
        self.assertIn("Unsupported dtype", run.stderr)
        self.assertNotIn("Traceback", run.stderr)

    def test_cli_tp_zero_fails_cleanly(self):
        cfg_path = _ROOT / "tests" / "data" / "minimal_dense_config.json"
        env = dict(os.environ)
        env["PYTHONPATH"] = str(_SRC)

        run = subprocess.run(
            [sys.executable, "-m", "llm_mem_planner", "--config", str(cfg_path), "--tp", "0"],
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=False,
        )
        self.assertNotEqual(run.returncode, 0)
        self.assertIn("tp >= 1", run.stderr)
        self.assertNotIn("Traceback", run.stderr)

    def test_cli_malformed_json_fails_cleanly(self):
        cfg_path = _ROOT / "tests" / "data" / "invalid_json_config.json"
        env = dict(os.environ)
        env["PYTHONPATH"] = str(_SRC)

        run = subprocess.run(
            [sys.executable, "-m", "llm_mem_planner", "--config", str(cfg_path), "--tp", "1"],
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=False,
        )
        self.assertNotEqual(run.returncode, 0)
        self.assertIn("Failed to parse JSON", run.stderr)
        self.assertNotIn("Traceback", run.stderr)

    def test_cli_missing_required_field_fails_cleanly(self):
        cfg_path = _ROOT / "tests" / "data" / "missing_hidden_size_config.json"
        env = dict(os.environ)
        env["PYTHONPATH"] = str(_SRC)

        run = subprocess.run(
            [sys.executable, "-m", "llm_mem_planner", "--config", str(cfg_path), "--tp", "1"],
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=False,
        )
        self.assertNotEqual(run.returncode, 0)
        self.assertIn("Config missing required field", run.stderr)
        self.assertIn("hidden_size", run.stderr)
        self.assertNotIn("Traceback", run.stderr)

    def test_cli_invalid_sp_candidates_fails_cleanly(self):
        cfg_path = _ROOT / "tests" / "data" / "minimal_dense_config.json"
        env = dict(os.environ)
        env["PYTHONPATH"] = str(_SRC)

        with tempfile.TemporaryDirectory(prefix="llm-mem-planner-test-") as td:
            tmp = Path(td)
            run = subprocess.run(
                [
                    sys.executable,
                    "-m",
                    "llm_mem_planner",
                    "--config",
                    str(cfg_path),
                    "--recommend",
                    "--nodes",
                    "1",
                    "--gpus-per-node",
                    "8",
                    "--gpu-mem-gib",
                    "80",
                    "--tp",
                    "8",
                    "--seq-len",
                    "2048",
                    "--batch",
                    "16",
                    "--sp-candidates",
                    "1,x",
                    "--out-md",
                    str(tmp / "out.md"),
                    "--out-svg",
                    str(tmp / "out.svg"),
                    "--out-perf-svg",
                    str(tmp / "out_perf.svg"),
                ],
                env=env,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                check=False,
            )
        self.assertNotEqual(run.returncode, 0)
        self.assertIn("--sp-candidates", run.stderr)
        self.assertNotIn("Traceback", run.stderr)

    def test_cli_recommend_without_tp_defaults_to_world_size(self):
        cfg_path = _ROOT / "tests" / "data" / "minimal_dense_config.json"
        env = dict(os.environ)
        env["PYTHONPATH"] = str(_SRC)

        with tempfile.TemporaryDirectory(prefix="llm-mem-planner-test-") as td:
            tmp = Path(td)
            run = subprocess.run(
                [
                    sys.executable,
                    "-m",
                    "llm_mem_planner",
                    "--config",
                    str(cfg_path),
                    "--recommend",
                    "--nodes",
                    "1",
                    "--gpus-per-node",
                    "8",
                    "--gpu-mem-gib",
                    "80",
                    "--seq-len",
                    "2048",
                    "--batch",
                    "16",
                    "--sp-candidates",
                    "1",
                    "--out-md",
                    str(tmp / "out.md"),
                    "--out-svg",
                    str(tmp / "out.svg"),
                    "--out-perf-svg",
                    str(tmp / "out_perf.svg"),
                ],
                env=env,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                check=False,
            )
        self.assertEqual(run.returncode, 0, msg=run.stderr)
        self.assertIn("== Recommendation Output ==", run.stdout)
        self.assertIn("Top feasible configs:", run.stdout)

    def test_cli_recommend_tp_zero_fails_cleanly(self):
        cfg_path = _ROOT / "tests" / "data" / "minimal_dense_config.json"
        env = dict(os.environ)
        env["PYTHONPATH"] = str(_SRC)

        with tempfile.TemporaryDirectory(prefix="llm-mem-planner-test-") as td:
            tmp = Path(td)
            run = subprocess.run(
                [
                    sys.executable,
                    "-m",
                    "llm_mem_planner",
                    "--config",
                    str(cfg_path),
                    "--recommend",
                    "--nodes",
                    "1",
                    "--gpus-per-node",
                    "8",
                    "--gpu-mem-gib",
                    "80",
                    "--tp",
                    "0",
                    "--seq-len",
                    "2048",
                    "--batch",
                    "16",
                    "--sp-candidates",
                    "1",
                    "--out-md",
                    str(tmp / "out.md"),
                    "--out-svg",
                    str(tmp / "out.svg"),
                    "--out-perf-svg",
                    str(tmp / "out_perf.svg"),
                ],
                env=env,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                check=False,
            )
        self.assertNotEqual(run.returncode, 0)
        self.assertIn("Invalid --tp", run.stderr)
        self.assertNotIn("Traceback", run.stderr)


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