import json
import unittest
import warnings
from types import SimpleNamespace

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_NIGHTLY_EVAL_FP8_TP1,
    DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP2,
    DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP1,
    DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP2,
    DEFAULT_URL_FOR_TEST,
    ModelLaunchSettings,
    check_evaluation_test_results,
    parse_models,
    popen_launch_server,
    write_results_to_json,
)

# Nightly eval tests run large models (up to 70B+ params) that may need
# downloading on cache miss. Use a longer timeout than the default 600s.
NIGHTLY_EVAL_SERVER_TIMEOUT = 1800

register_cuda_ci(est_time=3600, suite="nightly-eval-text-2-gpu", nightly=True)

MODEL_SCORE_THRESHOLDS = {
    # sgl-eval (zero-shot chat, \boxed{}, math_verify grading). Thresholds are
    # measured_score - 0.05, baselined on H100 2-GPU over the full 1319 split.
    "meta-llama/Llama-3.1-8B-Instruct": 0.77,  # 81.05% measured - 5%
    "Qwen/Qwen3-8B": 0.76,  # 81.43% measured - 5%
    "Qwen/Qwen3-4B": 0.77,  # 82.41% measured - 5%
    "meta-llama/Llama-3.1-70B-Instruct": 0.90,  # 94.77% measured - 5%
    "mistralai/Mixtral-8x7B-Instruct-v0.1": 0.39,  # 43.52% measured - 5%
    "Qwen/Qwen2-57B-A14B-Instruct": 0.46,  # 50.87% measured - 5%
    "neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8": 0.77,  # 82.34% measured - 5%
    "neuralmagic/Mistral-7B-Instruct-v0.3-FP8": 0.23,  # 27.82% measured - 5%
    "neuralmagic/DeepSeek-Coder-V2-Lite-Instruct-FP8": 0.80,  # 84.91% measured - 5%
    "zai-org/GLM-4.5-Air-FP8": 0.73,  # 77.48% measured - 5%
    "neuralmagic/gemma-2-2b-it-FP8": 0.02,  # 6.52% measured - 5%
    "neuralmagic/Meta-Llama-3.1-70B-Instruct-FP8": 0.89,  # 94.01% measured - 5%
    "neuralmagic/Mixtral-8x7B-Instruct-v0.1-FP8": 0.35,  # 40.33% measured - 5%
    "neuralmagic/Qwen2-72B-Instruct-FP8": 0.83,  # 87.64% measured - 5%
    "neuralmagic/Qwen2-57B-A14B-Instruct-FP8": 0.40,  # 44.66% measured - 5%
}


# Do not use `CustomTestCase` since `test_gsm8k_all_models` does not want retry
class TestNightlyGsm8KEval(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.models = []
        models_tp1 = parse_models(
            DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP1
        ) + parse_models(DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP1)
        for model_path in models_tp1:
            cls.models.append(ModelLaunchSettings(model_path, tp_size=1))

        models_tp2 = parse_models(
            DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP2
        ) + parse_models(DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP2)
        for model_path in models_tp2:
            cls.models.append(ModelLaunchSettings(model_path, tp_size=2))

        cls.base_url = DEFAULT_URL_FOR_TEST

    def test_gsm8k_all_models(self):
        warnings.filterwarnings(
            "ignore", category=ResourceWarning, message="unclosed.*socket"
        )
        is_first = True
        all_results = []
        for model_setup in self.models:
            with self.subTest(model=model_setup.model_path):
                other_args = list(model_setup.extra_args)
                process = None

                if model_setup.model_path == "meta-llama/Llama-3.1-70B-Instruct":
                    other_args.extend(["--mem-fraction-static", "0.9"])

                try:
                    process = popen_launch_server(
                        model=model_setup.model_path,
                        other_args=other_args,
                        base_url=self.base_url,
                        timeout=NIGHTLY_EVAL_SERVER_TIMEOUT,
                    )

                    args = SimpleNamespace(
                        base_url=self.base_url,
                        model=model_setup.model_path,
                        eval_name="gsm8k",
                        api="sgl_eval",
                        num_examples=None,
                        num_threads=1024,
                    )

                    metrics = run_eval(args)
                    print(
                        f"{'=' * 42}\n{model_setup.model_path} - metrics={metrics} score={metrics['score']}\n{'=' * 42}\n"
                    )

                    write_results_to_json(
                        model_setup.model_path, metrics, "w" if is_first else "a"
                    )
                    is_first = False

                    all_results.append(
                        (model_setup.model_path, metrics["score"], 0.0, None)
                    )
                except Exception as e:
                    error_message = str(e)
                    all_results.append(
                        (model_setup.model_path, None, None, error_message)
                    )
                    print(f"Error evaluating {model_setup.model_path}: {error_message}")
                finally:
                    if process is not None:
                        kill_process_tree(process.pid)

        try:
            with open("results.json", "r") as f:
                print("\nFinal Results from results.json:")
                print(json.dumps(json.load(f), indent=2))
        except Exception as e:
            print(f"Error reading results.json: {e}")

        # Check all scores after collecting all results
        check_evaluation_test_results(
            all_results,
            self.__class__.__name__,
            model_accuracy_thresholds=MODEL_SCORE_THRESHOLDS,
            model_count=len(self.models),
        )


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