"""Test suite for LLM implementations."""

import base64
import os

from unittest.mock import MagicMock, patch

import pytest

try:
    import lumen.ai as lmai

    from lumen.ai.agents.vega_lite import VegaLiteAgent
    from lumen.ai.llm import (
        Anthropic, AzureOpenAI, Google, Groq, Message, MistralAI, OpenAI,
    )

except ModuleNotFoundError:
    pytest.skip("lumen.ai could not be imported, skipping tests.", allow_module_level=True)

from instructor.processing.multimodal import Image

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_test_image() -> Image:
    """Create a tiny 1x1 PNG encoded as an instructor Image."""
    pixel = base64.b64encode(
        b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01'
        b'\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde\x00'
        b'\x00\x00\x0cIDATx\x9cc\xf8\x0f\x00\x00\x01\x01\x00'
        b'\x05\x18\xd8N\x00\x00\x00\x00IEND\xaeB`\x82'
    ).decode('utf-8')
    return Image.from_raw_base64(pixel)


# ---------------------------------------------------------------------------
# AzureOpenAI model kwargs tests
# ---------------------------------------------------------------------------

def test_api_key_env_var_defaults():
    """Each provider has the correct api_key_env_var class variable by default."""
    assert OpenAI.api_key_env_var == "OPENAI_API_KEY"
    assert Anthropic.api_key_env_var == "ANTHROPIC_API_KEY"
    assert Google.api_key_env_var == "GEMINI_API_KEY"
    assert MistralAI.api_key_env_var == "MISTRAL_API_KEY"


def test_api_key_populated_from_env_var(monkeypatch):
    """api_key is populated from api_key_env_var env var when not explicitly passed."""
    monkeypatch.setenv("OPENAI_API_KEY", "test-openai-key")
    llm = OpenAI(model_kwargs={"default": {"model": "gpt-4.1-mini"}})
    assert llm.api_key == "test-openai-key"


def test_explicit_api_key_takes_priority(monkeypatch):
    """Explicitly passed api_key takes priority over the env var."""
    monkeypatch.setenv("OPENAI_API_KEY", "env-key")
    llm = OpenAI(api_key="explicit-key", model_kwargs={"default": {"model": "gpt-4.1-mini"}})
    assert llm.api_key == "explicit-key"


def test_api_key_env_var_override_on_class(monkeypatch):
    """Overriding api_key_env_var on the class causes api_key to be read from the new env var."""
    monkeypatch.setenv("MY_CUSTOM_KEY", "custom-key-value")
    original = OpenAI.api_key_env_var
    try:
        OpenAI.api_key_env_var = "MY_CUSTOM_KEY"
        llm = OpenAI(model_kwargs={"default": {"model": "gpt-4.1-mini"}})
        assert llm.api_key == "custom-key-value"
    finally:
        OpenAI.api_key_env_var = original


def test_api_key_env_var_override_on_subclass(monkeypatch):
    """Overriding api_key_env_var on a subclass doesn't affect the parent class."""
    monkeypatch.setenv("MY_SUBCLASS_KEY", "subclass-key-value")

    class MyOpenAI(OpenAI):
        api_key_env_var = "MY_SUBCLASS_KEY"

    llm = MyOpenAI(model_kwargs={"default": {"model": "gpt-4.1-mini"}})
    assert llm.api_key == "subclass-key-value"
    assert OpenAI.api_key_env_var == "OPENAI_API_KEY"


def test_api_key_none_when_env_var_unset(monkeypatch):
    """api_key is None when api_key_env_var is set but the env var is not."""
    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
    llm = OpenAI(model_kwargs={"default": {"model": "gpt-4.1-mini"}})
    assert llm.api_key is None


def test_api_key_from_modified_provider_env_vars(monkeypatch):
    """Modifying PROVIDER_ENV_VARS and api_key_env_var is reflected at instantiation."""
    monkeypatch.setitem(lmai.llm.PROVIDER_ENV_VARS, "openai", "MY_PATCHED_KEY")
    monkeypatch.setenv("MY_PATCHED_KEY", "patched-key-value")
    original = OpenAI.api_key_env_var
    try:
        OpenAI.api_key_env_var = lmai.llm.PROVIDER_ENV_VARS["openai"]
        llm = OpenAI(model_kwargs={"default": {"model": "gpt-4.1-mini"}})
        assert llm.api_key == "patched-key-value"
    finally:
        OpenAI.api_key_env_var = original


def test_get_available_llm_returns_local_provider_when_no_env_vars_set(monkeypatch):
    """get_available_llm returns a local provider when no cloud API keys are set."""
    for env_var in lmai.llm.PROVIDER_ENV_VARS.values():
        monkeypatch.delenv(env_var, raising=False)
    # Local providers (no env var required) should still be returned
    result = lmai.llm.get_available_llm()
    assert result is not None


def test_get_available_llm_returns_correct_provider(monkeypatch):
    """get_available_llm returns the provider whose env var is set."""
    for env_var in lmai.llm.PROVIDER_ENV_VARS.values():
        monkeypatch.delenv(env_var, raising=False)
    monkeypatch.setenv("ANTHROPIC_API_KEY", "test-key")
    assert lmai.llm.get_available_llm() is Anthropic


def test_get_available_llm_respects_modified_provider_env_vars(monkeypatch):
    """get_available_llm picks up a modified PROVIDER_ENV_VARS entry."""
    for env_var in list(lmai.llm.PROVIDER_ENV_VARS.values()):
        monkeypatch.delenv(env_var, raising=False)
    monkeypatch.setitem(lmai.llm.PROVIDER_ENV_VARS, "openai", "MY_CUSTOM_OPENAI_KEY")
    monkeypatch.setenv("MY_CUSTOM_OPENAI_KEY", "custom-value")
    assert lmai.llm.get_available_llm() is OpenAI


async def test_azure_open_ai_get_model_kwargs():
    """Test that AzureOpenAI._get_model_kwargs merges instance config with model-specific kwargs."""
    model_kwargs = {
        "default": {"model": "d_m", "azure_ad_token_provider": "d_aatp"},
        "other": {"model": "r_m", "azure_ad_token_provider": "r_aatp"},
    }

    llm = AzureOpenAI(api_version="av", endpoint="ep", model_kwargs=model_kwargs)

    expected_default = {
        "model": "d_m",
        "azure_ad_token_provider": "d_aatp",
        "api_version": "av",
        "azure_endpoint": "ep",
    }
    assert llm._get_model_kwargs("default") == expected_default

    expected_other = {
        "model": "r_m",
        "azure_ad_token_provider": "r_aatp",
        "api_version": "av",
        "azure_endpoint": "ep",
    }
    assert llm._get_model_kwargs("other") == expected_other


async def test_azure_open_ai_get_model_kwargs_individual_models():
    """Test model-specific config overrides instance defaults.

    To support use case where models do not share api_version and endpoint.
    """
    model_kwargs = {
        "default": {
            "model": "d_m",
            "azure_ad_token_provider": "d_aatp",
            "api_version": "d_av",
            "azure_endpoint": "d_ep",
        },
        "other": {
            "model": "r_m",
            "azure_ad_token_provider": "r_aatp",
            "api_version": "r_av",
            "azure_endpoint": "r_ep",
        },
    }

    llm = AzureOpenAI(api_version="av", endpoint="ep", model_kwargs=model_kwargs)

    assert llm._get_model_kwargs("default") == model_kwargs["default"]
    assert llm._get_model_kwargs("other") == model_kwargs["other"]


def test_groq_registered_in_llm_providers():
    """Test that the Groq provider is registered in LLM_PROVIDERS."""
    assert "groq" in lmai.llm.LLM_PROVIDERS
    assert lmai.llm.LLM_PROVIDERS["groq"] == "Groq"


def test_groq_registered_in_provider_env_vars():
    """Test that the Groq provider env var is registered."""
    assert "groq" in lmai.llm.PROVIDER_ENV_VARS
    assert lmai.llm.PROVIDER_ENV_VARS["groq"] == "GROQ_API_KEY"


def test_groq_api_key_env_var():
    """Test that Groq has the correct api_key_env_var."""
    assert Groq.api_key_env_var == "GROQ_API_KEY"


def test_groq_defaults():
    """Test that Groq has the correct default endpoint and model."""
    groq = Groq(api_key="test-key")
    assert groq.endpoint == "https://api.groq.com/openai/v1"
    assert groq.model_kwargs["default"]["model"] == "llama-3.3-70b-versatile"


def test_get_available_llm_selects_groq(monkeypatch):
    """Test that get_available_llm() selects Groq when only GROQ_API_KEY is set."""
    for env_var in lmai.llm.PROVIDER_ENV_VARS.values():
        monkeypatch.delenv(env_var, raising=False)
    monkeypatch.setenv("GROQ_API_KEY", "test-key")
    assert lmai.llm.get_available_llm() is Groq
# ---------------------------------------------------------------------------
# _check_for_image tests
# ---------------------------------------------------------------------------

class TestCheckForImage:
    """Tests for Llm._check_for_image (detection + serialization)."""

    def test_plain_text_no_image(self, llm):
        """Plain text messages return False."""
        messages: list[Message] = [
            {"role": "user", "content": "Hello"},
            {"role": "assistant", "content": "Hi there"},
        ]
        _, contains = llm._check_for_image(messages)
        assert not contains

    def test_single_image_detected(self, llm):
        """A bare Image in content is detected and serialized."""
        img = _make_test_image()
        messages: list[Message] = [{"role": "user", "content": img}]
        result, contains = llm._check_for_image(messages)
        assert contains
        assert isinstance(result[0]["content"], Image)

    def test_image_in_list_detected(self, llm):
        """An Image inside a list content is detected."""
        img = _make_test_image()
        messages: list[Message] = [
            {"role": "user", "content": ["Describe this chart:", img]},
        ]
        _, contains = llm._check_for_image(messages)
        assert contains

    def test_list_preserves_text_and_image(self, llm):
        """Mixed [str, Image] content keeps both parts after serialization."""
        img = _make_test_image()
        messages: list[Message] = [
            {"role": "user", "content": ["Describe this chart:", img]},
        ]
        result, contains = llm._check_for_image(messages)
        assert contains
        content = result[0]["content"]
        assert isinstance(content, list)
        assert len(content) == 2
        assert content[0] == "Describe this chart:"
        assert isinstance(content[1], Image)

    def test_list_without_images(self, llm):
        """A list of plain strings returns False."""
        messages: list[Message] = [
            {"role": "user", "content": ["part one", "part two"]},
        ]
        _, contains = llm._check_for_image(messages)
        assert not contains

    def test_multiple_images_in_list(self, llm):
        """Multiple images in one message are all serialized."""
        img1 = _make_test_image()
        img2 = _make_test_image()
        messages: list[Message] = [
            {"role": "user", "content": ["Compare:", img1, img2]},
        ]
        result, contains = llm._check_for_image(messages)
        assert contains
        content = result[0]["content"]
        assert isinstance(content, list)
        assert len(content) == 3
        assert content[0] == "Compare:"
        assert isinstance(content[1], Image)
        assert isinstance(content[2], Image)


# ---------------------------------------------------------------------------
# _prepare_vision_messages tests
# ---------------------------------------------------------------------------

class TestPrepareVisionMessages:
    """Tests for VegaLiteAgent._prepare_vision_messages fallback behavior."""

    @pytest.fixture
    def agent(self, llm):
        import warnings
        agent = VegaLiteAgent.__new__(VegaLiteAgent)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", DeprecationWarning)
            warnings.simplefilter("ignore", PendingDeprecationWarning)
            agent.llm = llm
        return agent

    def test_no_vision_with_fallback(self, agent):
        """With text_fallback, appends text-only message when vision unsupported."""
        agent.llm._supports_vision = False
        msgs = [{"role": "user", "content": "hello"}]
        result = agent._prepare_vision_messages(msgs, None, "Annotate this")
        assert len(result) == 2
        assert result[1] == {"role": "user", "content": "Annotate this"}

    def test_no_editor_with_fallback(self, agent):
        """With text_fallback, appends text-only message when editor is None."""
        agent.llm._supports_vision = True
        msgs = [{"role": "user", "content": "hello"}]
        result = agent._prepare_vision_messages(msgs, None, "Annotate this")
        assert len(result) == 2
        assert result[1] == {"role": "user", "content": "Annotate this"}

    def test_image_export_fails_with_fallback(self, agent):
        """With text_fallback, appends text-only message when image export fails."""
        agent.llm._supports_vision = True
        mock_editor = MagicMock()
        mock_editor.__class__ = type("VegaLiteEditor", (), {})
        msgs = [{"role": "user", "content": "hello"}]
        with patch.object(agent, '_export_plot_image', return_value=None):
            with patch('lumen.ai.agents.vega_lite.VegaLiteEditor', mock_editor.__class__):
                result = agent._prepare_vision_messages(msgs, mock_editor, "Annotate this")
        assert len(result) == 2
        assert result[1] == {"role": "user", "content": "Annotate this"}

    def test_vision_success_returns_image_message(self, agent):
        """When vision succeeds, appends message with [content, Image]."""
        agent.llm._supports_vision = True
        mock_editor = MagicMock()
        mock_editor.__class__ = type("VegaLiteEditor", (), {})
        fake_png = (
            b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01'
            b'\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde\x00'
            b'\x00\x00\x0cIDATx\x9cc\xf8\x0f\x00\x00\x01\x01\x00'
            b'\x05\x18\xd8N\x00\x00\x00\x00IEND\xaeB`\x82'
        )
        msgs = [{"role": "user", "content": "hello"}]
        with patch.object(agent, '_export_plot_image', return_value=fake_png):
            with patch('lumen.ai.agents.vega_lite.VegaLiteEditor', mock_editor.__class__):
                result = agent._prepare_vision_messages(msgs, mock_editor, "Annotate this")
        assert len(result) == 2
        content = result[1]["content"]
        assert isinstance(content, list)
        assert len(content) == 2
        assert content[0] == "Annotate this"
        assert isinstance(content[1], Image)
