from logging import getLogger
from unittest import IsolatedAsyncioTestCase, TestCase, main
from unittest.mock import AsyncMock, MagicMock, patch

from avalan.entities import (
    GenerationSettings,
    Message,
    MessageRole,
    TransformerEngineSettings,
)
from avalan.model import TextGenerationResponse
from avalan.model.nlp.text.vendor import TextGenerationVendorModel


class DummyVendorModel(TextGenerationVendorModel):
    pass


# allow instantiation without implementing abstract methods
DummyVendorModel.__abstractmethods__ = set()


class ConstructorTestCase(TestCase):
    def test_no_settings_requires_token(self):
        with self.assertRaises(AssertionError):
            DummyVendorModel("m", logger=getLogger())

    def test_settings_without_token_raises(self):
        settings = TransformerEngineSettings(
            auto_load_model=False, auto_load_tokenizer=False
        )
        with self.assertRaises(AssertionError):
            DummyVendorModel("m", settings, logger=getLogger())

    def test_settings_with_token_success(self):
        settings = TransformerEngineSettings(
            auto_load_model=False,
            auto_load_tokenizer=False,
            access_token="tok",
            enable_eval=True,
        )
        model = DummyVendorModel("m", settings, logger=getLogger())
        self.assertIsInstance(model, DummyVendorModel)
        self.assertFalse(model._settings.enable_eval)


class PropertyTestCase(TestCase):
    def setUp(self):
        self.settings = TransformerEngineSettings(
            auto_load_model=False,
            auto_load_tokenizer=False,
            access_token="tok",
        )
        self.model = DummyVendorModel("m", self.settings, logger=getLogger())

    def test_properties(self):
        self.assertFalse(self.model.supports_sample_generation)
        self.assertTrue(self.model.supports_token_streaming)
        self.assertFalse(self.model.uses_tokenizer)

    def test_load_model_not_implemented(self):
        with self.assertRaises(NotImplementedError):
            self.model._load_model()

    def test_tokenize_input_not_implemented(self):
        with self.assertRaises(NotImplementedError):
            self.model._tokenize_input("in")


class InputTokenCountTestCase(TestCase):
    def setUp(self):
        self.settings = TransformerEngineSettings(
            auto_load_model=False,
            auto_load_tokenizer=False,
            access_token="tok",
        )
        self.model = DummyVendorModel("m", self.settings, logger=getLogger())

    def test_resolve_encoding_with_model_specific_encoding(self):
        tiktoken = MagicMock()
        encoding = MagicMock()
        tiktoken.encoding_for_model.return_value = encoding

        with patch(
            "avalan.model.nlp.text.vendor.import_module",
            return_value=tiktoken,
        ) as import_module_mock:
            resolved = self.model._resolve_encoding("m", "cl100k_base")

        import_module_mock.assert_called_once_with("tiktoken")
        tiktoken.encoding_for_model.assert_called_once_with("m")
        tiktoken.get_encoding.assert_not_called()
        self.assertIs(resolved, encoding)

    def test_resolve_encoding_falls_back_to_default_encoding(self):
        tiktoken = MagicMock()
        encoding = MagicMock()
        tiktoken.encoding_for_model.side_effect = KeyError("missing")
        tiktoken.get_encoding.return_value = encoding

        with patch(
            "avalan.model.nlp.text.vendor.import_module",
            return_value=tiktoken,
        ):
            resolved = self.model._resolve_encoding("m", "cl100k_base")

        tiktoken.encoding_for_model.assert_called_once_with("m")
        tiktoken.get_encoding.assert_called_once_with("cl100k_base")
        self.assertIs(resolved, encoding)

    def test_input_token_count_with_model_encoding(self):
        encoding = MagicMock()
        encoding.encode.side_effect = lambda text: list(text)
        messages = [
            Message(role=MessageRole.USER, content="hi"),
            Message(role=MessageRole.USER, content="there"),
        ]
        with patch.object(
            self.model, "_resolve_encoding", return_value=encoding
        ) as resolve:
            self.model._messages = MagicMock(return_value=messages)
            count = self.model.input_token_count("in")
        resolve.assert_called_once_with("m", "cl100k_base")
        self.assertEqual(count, len("hi") + len("there"))

    def test_input_token_count_includes_instructions(self):
        encoding = MagicMock()
        encoding.encode.side_effect = lambda text: list(text)
        messages = [Message(role=MessageRole.USER, content="hi")]
        with patch.object(
            self.model, "_resolve_encoding", return_value=encoding
        ):
            self.model._messages = MagicMock(return_value=messages)
            count = self.model.input_token_count("in", instructions="provider")

        self.assertEqual(count, len("provider") + len("hi"))

    def test_input_token_count_fallback_encoding(self):
        encoding = MagicMock()
        encoding.encode.side_effect = lambda text: list(text)
        messages = [Message(role=MessageRole.USER, content="a")]
        with patch.object(
            self.model, "_resolve_encoding", return_value=encoding
        ) as resolve:
            self.model._messages = MagicMock(return_value=messages)
            count = self.model.input_token_count("in")
        resolve.assert_called_once_with("m", "cl100k_base")
        self.assertEqual(count, 1)

    def test_input_token_count_without_model_id_uses_default_encoding(self):
        model = DummyVendorModel("", self.settings, logger=getLogger())
        encoding = MagicMock()
        encoding.encode.side_effect = lambda text: list(text)
        messages = [Message(role=MessageRole.USER, content="abc")]
        with patch.object(
            model, "_resolve_encoding", return_value=encoding
        ) as resolve:
            model._messages = MagicMock(return_value=messages)
            count = model.input_token_count("in")
        resolve.assert_called_once_with("cl100k_base", "cl100k_base")
        self.assertEqual(count, 3)


class CallTestCase(IsolatedAsyncioTestCase):
    async def test_call(self):
        settings = TransformerEngineSettings(
            auto_load_model=False,
            auto_load_tokenizer=False,
            access_token="tok",
        )
        model = DummyVendorModel("m", settings, logger=getLogger())
        messages = [Message(role=MessageRole.USER, content="hi")]
        model._messages = MagicMock(return_value=messages)

        async def fake_model(
            model_id, msgs, opts, *, tool=None, use_async_generator=True
        ):
            self.assertEqual(model_id, "m")
            self.assertIs(msgs, messages)
            self.assertIs(opts, gen_settings)
            self.assertTrue(use_async_generator)
            return "streamer"

        model._model = AsyncMock(side_effect=fake_model)
        gen_settings = GenerationSettings()
        response = await model(
            "input",
            system_prompt="sys",
            settings=gen_settings,
            tool=None,
        )
        model._messages.assert_called_once_with("input", "sys", None, None)
        model._model.assert_awaited_once_with(
            "m",
            messages,
            gen_settings,
            tool=None,
            use_async_generator=True,
        )
        self.assertIsInstance(response, TextGenerationResponse)
        self.assertIs(response._output_fn, "streamer")
        self.assertIs(response._kwargs["settings"], gen_settings)
        self.assertTrue(response._use_async_generator)


if __name__ == "__main__":
    main()
