import unittest
from argparse import Namespace
from unittest.mock import MagicMock

from avalan.cli.commands import get_model_settings
from avalan.entities import Modality


class GetModelSettingsTestCase(unittest.TestCase):
    def test_basic(self):
        engine_uri = MagicMock()
        args = Namespace(
            attention="flash",
            base_url="http://localhost:9001/v1",
            device="cpu",
            parallel="colwise",
            disable_loading_progress_bar=True,
            sentence_transformer=True,
            loader_class="auto",
            backend="transformers",
            low_cpu_mem_usage=True,
            quiet=False,
            revision="rev",
            special_token=["<s>"],
            tokenizer="tok",
            token=["t"],
            subfolder="model_sub",
            tokenizer_subfolder="tok_sub",
            trust_remote_code=True,
            weight_type="fp16",
            output_hidden_states=False,
        )

        result = get_model_settings(args, MagicMock(), MagicMock(), engine_uri)
        expected = {
            "engine_uri": engine_uri,
            "attention": "flash",
            "output_hidden_states": False,
            "base_url": "http://localhost:9001/v1",
            "device": "cpu",
            "parallel": "colwise",
            "disable_loading_progress_bar": True,
            "modality": Modality.EMBEDDING,
            "loader_class": "auto",
            "backend": "transformers",
            "low_cpu_mem_usage": True,
            "quiet": False,
            "revision": "rev",
            "base_model_id": None,
            "checkpoint": None,
            "special_tokens": ["<s>"],
            "tokenizer": "tok",
            "tokens": ["t"],
            "subfolder": "model_sub",
            "tokenizer_subfolder": "tok_sub",
            "refiner_model_id": None,
            "upsampler_model_id": None,
            "trust_remote_code": True,
            "weight_type": "fp16",
        }
        self.assertEqual(result, expected)

    def test_override_sentence_transformer(self):
        engine_uri = MagicMock()
        args = Namespace(
            device="cpu",
            disable_loading_progress_bar=False,
            loader_class="auto",
            backend="transformers",
            low_cpu_mem_usage=False,
            quiet=True,
            revision="main",
            special_token=None,
            tokenizer=None,
            token=None,
            weight_type="fp32",
        )

        result = get_model_settings(
            args,
            MagicMock(),
            MagicMock(),
            engine_uri,
            modality=Modality.EMBEDDING,
        )
        self.assertEqual(result["modality"], Modality.EMBEDDING)

    def test_string_modality_is_normalized(self):
        engine_uri = MagicMock()
        engine_uri.params = {}
        args = Namespace(
            modality="text_generation",
            device="cpu",
            disable_loading_progress_bar=False,
            loader_class="auto",
            backend="mlx",
            low_cpu_mem_usage=False,
            quiet=True,
            revision=None,
            special_token=None,
            tokenizer=None,
            token=None,
            weight_type="bf16",
        )

        result = get_model_settings(args, MagicMock(), MagicMock(), engine_uri)

        self.assertEqual(result["modality"], Modality.TEXT_GENERATION)

    def test_ds4_cli_options_create_backend_config(self):
        engine_uri = MagicMock()
        engine_uri.params = {}
        args = Namespace(
            device="cpu",
            disable_loading_progress_bar=False,
            loader_class="auto",
            backend="ds4",
            low_cpu_mem_usage=False,
            quiet=True,
            revision=None,
            special_token=None,
            tokenizer=None,
            token=None,
            weight_type="bf16",
            ds4_ctx=4096,
            ds4_native_backend="metal",
            ds4_mtp="mtp.gguf",
            ds4_mtp_draft=2,
            ds4_mtp_margin=0.1,
            ds4_warm_weights=True,
            ds4_quality=True,
            ds4_native_log=True,
        )

        result = get_model_settings(args, MagicMock(), MagicMock(), engine_uri)

        self.assertEqual(
            result["backend_config"],
            {
                "ctx_size": 4096,
                "native_backend": "metal",
                "mtp_path": "mtp.gguf",
                "mtp_draft_tokens": 2,
                "mtp_margin": 0.1,
                "warm_weights": True,
                "quality": True,
                "native_log": True,
            },
        )

    def test_ds4_cli_omits_native_log_when_unspecified(self):
        engine_uri = MagicMock()
        engine_uri.params = {}
        args = Namespace(
            device="cpu",
            disable_loading_progress_bar=False,
            loader_class="auto",
            backend="ds4",
            low_cpu_mem_usage=False,
            quiet=True,
            revision=None,
            special_token=None,
            tokenizer=None,
            token=None,
            weight_type="bf16",
            ds4_native_log=None,
        )

        result = get_model_settings(args, MagicMock(), MagicMock(), engine_uri)

        self.assertNotIn("backend_config", result)

    def test_ds4_cli_options_do_not_affect_non_ds4_backend(self):
        engine_uri = MagicMock()
        engine_uri.params = {}
        args = Namespace(
            device="cpu",
            disable_loading_progress_bar=False,
            loader_class="auto",
            backend="transformers",
            low_cpu_mem_usage=False,
            quiet=True,
            revision=None,
            special_token=None,
            tokenizer=None,
            token=None,
            weight_type="bf16",
            ds4_ctx=4096,
        )

        result = get_model_settings(args, MagicMock(), MagicMock(), engine_uri)

        self.assertNotIn("backend_config", result)
