import os
import unittest
from pathlib import Path
from unittest.mock import patch

from core.capabilities import ModelCapabilities
from core.llm import LLM


class LLMKwargsTests(unittest.TestCase):
    def _llm(
        self, *, family: str, thinking_enabled: bool, thinking_transport: str
    ) -> LLM:
        caps = ModelCapabilities(
            family=family,
            model_id=f"{family}/model",
            api_key_env="TEST_LLM_API_KEY",
            thinking_enabled=thinking_enabled,
            thinking_transport=thinking_transport,
            optimal_temperature=0.3,
        )
        with patch.dict(os.environ, {"TEST_LLM_API_KEY": "test-key"}):
            return LLM(caps)

    def test_deepseek_explicitly_enables_thinking_and_sets_effort(self) -> None:
        llm = self._llm(
            family="deepseek_v4", thinking_enabled=True, thinking_transport="extra_body"
        )

        kwargs = llm._build_kwargs(
            [{"role": "user", "content": "hello"}], None, None, "high"
        )

        self.assertNotIn("reasoning_effort", kwargs)
        self.assertEqual(
            kwargs["extra_body"],
            {"thinking": {"type": "enabled"}, "reasoning_effort": "high"},
        )

    def test_deepseek_explicitly_disables_thinking_without_effort(self) -> None:
        llm = self._llm(
            family="deepseek_v4", thinking_enabled=False, thinking_transport="extra_body"
        )

        kwargs = llm._build_kwargs(
            [{"role": "user", "content": "hello"}], None, None, "high"
        )

        self.assertNotIn("reasoning_effort", kwargs)
        self.assertEqual(
            kwargs["extra_body"], {"thinking": {"type": "disabled"}}
        )

    def test_other_openai_compatible_provider_gets_no_thinking_body(self) -> None:
        llm = self._llm(
            family="unifyllm", thinking_enabled=False, thinking_transport="none"
        )

        kwargs = llm._build_kwargs(
            [{"role": "user", "content": "hello"}], None, None, None
        )

        self.assertNotIn("extra_body", kwargs)

    def test_extra_body_transport_enables_thinking_without_effort(self) -> None:
        llm = self._llm(
            family="doubao", thinking_enabled=True, thinking_transport="extra_body"
        )

        kwargs = llm._build_kwargs(
            [{"role": "user", "content": "hello"}], None, None, None
        )

        self.assertEqual(
            kwargs["extra_body"], {"thinking": {"type": "enabled"}}
        )

    def test_invalid_transport_is_rejected_before_request(self) -> None:
        llm = self._llm(
            family="test", thinking_enabled=True, thinking_transport="unknown"
        )

        with self.assertRaisesRegex(ValueError, "thinking_transport"):
            llm._build_kwargs(
                [{"role": "user", "content": "hello"}], None, None, "high"
            )

    def test_flash_profile_matches_0731_capabilities(self) -> None:
        caps = ModelCapabilities.load(
            "deepseek_v4.flash", Path(__file__).resolve().parents[1] / "config" / "models"
        )

        self.assertTrue(caps.thinking_enabled)
        self.assertEqual(caps.reasoning_effort_levels, ["low", "high", "max"])
        self.assertEqual(caps.default_reasoning_effort, "high")
        self.assertEqual(caps.max_output, 8192)
        self.assertEqual(caps.output_cny_per_mtoken, 2.0)
        self.assertEqual(caps.cache_hit_cny_per_mtoken, 0.02)
        self.assertEqual(caps.thinking_transport, "extra_body")

    def test_other_controllable_profiles_declare_transport(self) -> None:
        models_dir = Path(__file__).resolve().parents[1] / "config" / "models"

        glm = ModelCapabilities.load("glm.pro52", models_dir)
        doubao = ModelCapabilities.load("doubao.turbo", models_dir)

        self.assertFalse(glm.thinking_enabled)
        self.assertEqual(glm.thinking_transport, "extra_body")
        self.assertTrue(doubao.thinking_enabled)
        self.assertEqual(doubao.thinking_transport, "extra_body")
        self.assertEqual(doubao.default_reasoning_effort, "")


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