from __future__ import annotations

import re
import sys
import unittest
from pathlib import Path


PACKAGE_DIR = Path(__file__).resolve().parents[1] / "modelhub_submmit_api"
if str(PACKAGE_DIR) in sys.path:
    sys.path.remove(str(PACKAGE_DIR))
sys.path.insert(0, str(PACKAGE_DIR))

from template_selector import TemplateSelector  # noqa: E402


class PlatformTemplateValidationTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.selector = TemplateSelector()

    def test_removed_kunlunxin_gpu_is_not_auto_selected(self) -> None:
        supported = self.selector.supported_target_gpus("text-generation", auto_only=True)
        self.assertNotIn("Kunlunxin_r-200-8f", supported)
        self.assertEqual(
            set(),
            self.selector.supported_frameworks_for_auto("text-generation", "Kunlunxin_r-200-8f"),
        )
        with self.assertRaisesRegex(KeyError, "no longer accepted"):
            self.selector.normalize_gpu("Kunlunxin_r-200-8f")

    def test_biren_template_is_rendered_with_single_gpu_parallelism(self) -> None:
        template = self.selector.select_template("text-generation", "vllm", "Biren_166m")
        config = self.selector.render_config(template)

        gpu_counts = re.findall(r"(?m)^\s*gpu_num:\s*(\d+)", config)
        tensor_parallel_values = re.findall(r"(?:-tp|--tensor-parallel-size)(?:\s+|=|,\s*['\"]?)(\d+)", config)
        self.assertEqual(["1", "1"], gpu_counts)
        self.assertEqual(["1", "1"], tensor_parallel_values)
        self.assertNotIn("gpu_num: 4", config)


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