"""
Usage:
python3 -m unittest test_autoround_quantization
"""

import os
import shutil
import tempfile
import unittest
from types import SimpleNamespace

from sglang.srt.configs.device_config import DeviceConfig
from sglang.srt.configs.load_config import LoadConfig
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.model_loader.loader import get_model_loader
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.run_eval import run_eval
from sglang.test.test_utils import (
    DEFAULT_MODEL_NAME_FOR_TEST,
    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
    DEFAULT_URL_FOR_TEST,
    CustomTestCase,
    popen_launch_server,
)

register_cuda_ci(est_time=220, stage="extra-a", runner_config="1-gpu-large")

MMLU_NUM_EXAMPLES = 256
MMLU_NUM_THREADS = 32
# The unquantized model scores 0.68-0.70 over 256 questions; int8 gives up a
# point or two, then the 0.05 margin the other eval thresholds use.
MMLU_SCORE_THRESHOLD = 0.63


class TestAutoRoundQuantization(CustomTestCase):
    @classmethod
    def setUpClass(cls):
        cls.base_url = DEFAULT_URL_FOR_TEST
        cls.model = DEFAULT_MODEL_NAME_FOR_TEST
        cls.output_dir = tempfile.mkdtemp()

    @classmethod
    def tearDownClass(cls):
        if os.path.isdir(cls.output_dir):
            shutil.rmtree(cls.output_dir)

    def test_online_quant(self):
        process = popen_launch_server(
            self.model,
            self.base_url,
            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
            other_args=["--trust-remote-code", "--quantization", "auto-round-int8"],
        )

        try:
            args = SimpleNamespace(
                base_url=self.base_url,
                model=self.model,
                eval_name="mmlu",
                num_examples=MMLU_NUM_EXAMPLES,
                num_threads=MMLU_NUM_THREADS,
                device="auto",
            )
            metrics = run_eval(args)
            self.assertGreaterEqual(metrics["score"], MMLU_SCORE_THRESHOLD)
        finally:
            kill_process_tree(process.pid)
            print(f"[INFO] Server for {self.model} stopped.")

    def test_offline_quant(self):
        model_config = ModelConfig(
            model_path=self.model,
            quantization="auto-round-int8",
            trust_remote_code=True,
        )
        load_config = LoadConfig(
            inc_save_path=self.output_dir,
        )
        device_config = DeviceConfig(device="cuda")
        model_loader = get_model_loader(load_config, model_config)
        quantized_model = model_loader.load_model(
            model_config=model_config,
            device_config=device_config,
        )
        # AutoRound saves the quantized checkpoint into a scheme-derived
        # subfolder (e.g. "<model>-w8a8/") under the output dir
        config_found = any(
            "config.json" in files for _, _, files in os.walk(self.output_dir)
        )
        self.assertTrue(
            config_found,
            f"No config.json written under {self.output_dir}",
        )


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