import contextlib
import io
import json
import os
import unittest
from unittest import mock

from prain.cli import main
from prain.config import load_config
from tests.helpers import TmpCase

GEMMA4 = {"text_config": {"num_hidden_layers": 30, "num_experts": 128, "top_k_experts": 8,
                          "hidden_size": 2816, "moe_intermediate_size": 704}}


class Cli(TmpCase):
    def run_cli(self, *argv):
        out, err = io.StringIO(), io.StringIO()
        with contextlib.redirect_stdout(out), contextlib.redirect_stderr(err):
            main(list(argv))
        return out.getvalue(), err.getvalue()

    def test_check_prints_tier_sizes(self):
        hf = self.write("config.json", json.dumps(GEMMA4))
        config = self.write("p.toml", 'pack_dir = "/packs"\ncache_slots = 4\n')
        out, err = self.run_cli("check", config, hf)
        self.assertIn("    0     0   128     0            4", out)
        self.assertTrue(out.endswith(
            "expert size 11.34 MiB\n"
            "VRAM 1.33 GiB  (pinned experts + cache slots)\n"
            "KV   0.75 GiB  (shared by up to 2 parallel requests)\n"
            "RAM  42.54 GiB  (pinned host memory)\n"
            "disk 42.54 GiB  (pack in /packs)\n"), out)
        self.assertIn("warning: 30 layers have fewer cache slots than top_k=8", err)

    def test_plan_output_is_a_valid_config(self):
        hf = self.write("config.json", json.dumps(GEMMA4))
        stats = {"layers": {str(layer): {"uses": list(range(128))} for layer in range(30)}}
        stats_path = self.write("stats.json", json.dumps(stats))
        out, _ = self.run_cli("plan", stats_path, hf, "--expert-vram-gib", "8",
                              "--ram-gib", "16", "--pack-dir", "/packs")
        config = self.write("planned.toml", out)
        check_out, _ = self.run_cli("check", config, hf)
        self.assertIn("VRAM 8.00 GiB", check_out)  # budgets are filled, never exceeded
        self.assertIn("RAM  16.00 GiB", check_out)

    def test_serve_execs_vllm_with_the_required_flags(self):
        config = self.write("p.toml", 'pack_dir = "/packs"\n[miss]\nmode = "auto"\ncpu_threads = 6\n'
                                      '[serve]\nmax_requests = 3\nkv_cache_gib = 1.0\n')
        with mock.patch("os.execvp") as execvp, mock.patch.dict(os.environ, clear=False):
            os.environ.pop("OMP_NUM_THREADS", None)
            main(["serve", config, "/models/m", "--", "--port", "8001"])
            self.assertEqual(os.environ["PRAIN_CONFIG"], config)
            self.assertEqual(os.environ["OMP_NUM_THREADS"], "6")
        execvp.assert_called_once_with("vllm", [
            "vllm", "serve", "/models/m", "--enforce-eager", "--moe-backend", "triton",
            "--scheduling-policy", "priority", "--max-num-seqs", "3",
            "--kv-cache-memory-bytes", "1073741824", "--chat-template-content-format", "string",
            "--port", "8001"])
        with mock.patch("os.execvp") as execvp, self.assertRaisesRegex(
                SystemExit, "--max-num-seqs=4 comes from the placement file"):
            main(["serve", config, "/models/m", "--", "--max-num-seqs=4"])
        execvp.assert_not_called()

    def test_web_binds_the_host_it_is_given(self):
        with mock.patch("prain.web.serve") as web:
            main(["web", "--port", "8081"])
            self.assertEqual(web.call_args.kwargs["host"], "127.0.0.1")
            main(["web", "--host", "0.0.0.0"])
            self.assertEqual(web.call_args.kwargs["host"], "0.0.0.0")

    def test_a_persona_file_reaches_web_and_tui(self):
        path = self.write("gentoo-admin.md", "You look after a Gentoo box.\n\n")
        with mock.patch("prain.web.serve") as web:
            main(["web", "--persona", path])
            self.assertEqual((web.call_args.kwargs["system"], web.call_args.kwargs["persona"]),
                             ("You look after a Gentoo box.", "gentoo-admin"))
            main(["web"])
            self.assertEqual((web.call_args.kwargs["system"], web.call_args.kwargs["persona"]),
                             (None, None))
        with mock.patch("prain.tui.main") as tui:
            main(["tui", "--persona", path])
            self.assertEqual(tui.call_args.kwargs["system"], "You look after a Gentoo box.")

    def test_synth_sim_and_tune(self):
        hf = self.write("config.json", json.dumps(GEMMA4))
        trace = os.path.join(self.dir, "trace.jsonl.gz")
        self.run_cli("synth", trace, "--forwards", "40", "--layers", "3")
        config = self.write("p.toml", 'pack_dir = "/packs"\ncache_slots = 5\n')
        out, _ = self.run_cli("sim", trace, config, hf)
        self.assertRegex(out, r"VRAM hit rate \d+\.\d%  \(eviction alone at best: \d+\.\d%\)")
        self.assertIn("decode ", out)
        out, err = self.run_cli("tune", trace, hf, "--expert-vram-gib", "0.2",
                                "--ram-gib", "4", "--pack-dir", "/packs", "--top", "3")
        self.assertIn("    1  ", err)
        tuned = load_config(self.write("tuned.toml", out))
        self.assertEqual((tuned.pack_dir, tuned.trace_path), ("/packs", trace))
        self.assertLessEqual(sum(tuned.placement(layer, 128).cache_slots for layer in range(3)),
                             int(0.2 * 1024 / 11.34) // 3 * 3)

    def test_quantized_experts_check_and_tune(self):
        hf = self.write("config.json", json.dumps(GEMMA4))
        out, _ = self.run_cli("check", self.write("q.toml", 'pack_dir = "/packs"\nquant = "int4"\n'), hf)
        self.assertIn("expert size 3.01 MiB\n", out)
        trace = os.path.join(self.dir, "trace.jsonl.gz")
        self.run_cli("synth", trace, "--forwards", "40", "--layers", "3")
        out, _ = self.run_cli("tune", trace, hf, "--quant", "int4", "--expert-vram-gib", "0.2",
                              "--ram-gib", "4", "--pack-dir", "/packs", "--top", "3")
        tuned = load_config(self.write("tuned.toml", out))
        self.assertEqual(tuned.quant, "int4")
        slots = sum(tuned.placement(layer, 128).cache_slots for layer in range(3))
        self.assertGreater(slots, int(0.2 * 1024 / 11.34))  # int4 experts are ~4x smaller
        self.assertLessEqual(slots, int(0.2 * 1024 / 3.01))