from unittest.mock import MagicMock, patch

import pytest

from mem0.configs.llms.aws_bedrock import AWSBedrockConfig
from mem0.llms.aws_bedrock import AWSBedrockLLM, extract_provider
from mem0.utils.factory import LlmFactory

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

@pytest.fixture
def mock_boto3():
    """Patch boto3 so no real AWS calls are made during unit tests."""
    with patch("mem0.llms.aws_bedrock.boto3") as mock_b3:
        runtime_client = MagicMock()
        bedrock_client = MagicMock()
        bedrock_client.list_foundation_models.return_value = {"modelSummaries": []}

        def _client(service, **kwargs):
            if service == "bedrock-runtime":
                return runtime_client
            return bedrock_client

        mock_b3.client.side_effect = _client
        yield runtime_client


def _make_llm(model: str, mock_boto3, **kwargs) -> AWSBedrockLLM:
    """Instantiate AWSBedrockLLM with a given model, all AWS calls mocked."""
    config = AWSBedrockConfig(model=model, **kwargs)
    return AWSBedrockLLM(config)


def _converse_response(text: str = "ok") -> dict:
    """Minimal Converse API response dict."""
    return {"output": {"message": {"content": [{"text": text}]}}}


# ---------------------------------------------------------------------------
# extract_provider
# ---------------------------------------------------------------------------

class TestExtractProvider:
    def test_standard_anthropic_model(self):
        assert extract_provider("anthropic.claude-3-5-sonnet-20240620-v1:0") == "anthropic"

    def test_inference_profile_us_prefix(self):
        # Cross-region inference profile IDs look like us.anthropic.<model>
        assert extract_provider("us.anthropic.claude-haiku-4-5-20251001-v1:0") == "anthropic"

    def test_inference_profile_eu_prefix(self):
        assert extract_provider("eu.anthropic.claude-sonnet-4-5-20250929-v1:0") == "anthropic"

    def test_inference_profile_ap_prefix(self):
        assert extract_provider("ap.anthropic.claude-3-opus-20240229-v1:0") == "anthropic"

    def test_amazon_model(self):
        assert extract_provider("amazon.nova-3-mini-20241119-v1:0") == "amazon"

    def test_meta_model(self):
        assert extract_provider("meta.llama3-8b-instruct-v1:0") == "meta"

    def test_mistral_model(self):
        assert extract_provider("mistral.mistral-7b-instruct-v0:2") == "mistral"

    def test_unknown_model_raises(self):
        with pytest.raises(ValueError, match="Unknown provider"):
            extract_provider("unknown-vendor.some-model-v1:0")

    def test_application_inference_profile_arn_without_override_raises(self):
        arn = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123xyz"
        with pytest.raises(ValueError, match="Unknown provider"):
            extract_provider(arn)

    def test_application_inference_profile_arn_with_explicit_provider(self):
        arn = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123xyz"
        assert extract_provider(arn, "anthropic") == "anthropic"

    def test_explicit_provider_takes_precedence_over_regex(self):
        assert extract_provider("anthropic.claude-3-5-sonnet-20240620-v1:0", "amazon") == "amazon"

    def test_explicit_provider_typo_raises(self):
        arn = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123xyz"
        with pytest.raises(ValueError, match="Unknown provider_override 'anthorpic'"):
            extract_provider(arn, "anthorpic")


# ---------------------------------------------------------------------------
# AWSBedrockConfig
# ---------------------------------------------------------------------------

class TestAWSBedrockConfig:
    def test_top_p_defaults_to_none(self):
        config = AWSBedrockConfig(model="anthropic.claude-3-5-sonnet-20240620-v1:0")
        assert config.top_p is None

    def test_top_p_explicit_value_stored(self):
        config = AWSBedrockConfig(model="anthropic.claude-3-5-sonnet-20240620-v1:0", top_p=0.8)
        assert config.top_p == 0.8

    def test_get_model_config_excludes_top_p_by_default(self):
        config = AWSBedrockConfig(model="anthropic.claude-3-5-sonnet-20240620-v1:0", temperature=0.5)
        model_cfg = config.get_model_config()
        assert "top_p" not in model_cfg

    def test_get_model_config_includes_top_p_when_set(self):
        config = AWSBedrockConfig(model="anthropic.claude-3-5-sonnet-20240620-v1:0", top_p=0.7)
        model_cfg = config.get_model_config()
        assert model_cfg["top_p"] == 0.7

    def test_get_model_config_top_p_via_model_kwargs(self):
        """model_kwargs can supply top_p after merge; same semantics as explicit top_p."""
        config = AWSBedrockConfig(
            model="anthropic.claude-3-5-sonnet-20240620-v1:0",
            model_kwargs={"top_p": 0.88},
        )
        assert config.get_model_config()["top_p"] == 0.88

    def test_get_model_config_always_includes_temperature(self):
        config = AWSBedrockConfig(model="anthropic.claude-3-5-sonnet-20240620-v1:0", temperature=0.3)
        model_cfg = config.get_model_config()
        assert model_cfg["temperature"] == 0.3

    def test_aws_region_stored(self):
        config = AWSBedrockConfig(
            model="anthropic.claude-3-5-sonnet-20240620-v1:0",
            aws_region="us-east-2",
        )
        assert config.aws_region == "us-east-2"

    def test_provider_override_defaults_to_none(self):
        config = AWSBedrockConfig(model="anthropic.claude-3-5-sonnet-20240620-v1:0")
        assert config.provider_override is None

    def test_provider_override_stored(self):
        arn = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123xyz"
        config = AWSBedrockConfig(model=arn, provider_override="anthropic")
        assert config.provider_override == "anthropic"


# ---------------------------------------------------------------------------
# AWSBedrockLLM with application inference profile ARNs
# ---------------------------------------------------------------------------

class TestApplicationInferenceProfileArn:
    ARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123xyz"

    def test_arn_without_provider_override_raises(self, mock_boto3):
        with pytest.raises(ValueError, match="Unknown provider"):
            _make_llm(self.ARN, mock_boto3)

    def test_arn_with_provider_override_resolves(self, mock_boto3):
        llm = _make_llm(self.ARN, mock_boto3, provider_override="anthropic")
        assert llm.provider == "anthropic"
        assert llm.supports_tools is True

    def test_plain_model_id_unaffected(self, mock_boto3):
        llm = _make_llm("anthropic.claude-3-5-sonnet-20240620-v1:0", mock_boto3)
        assert llm.provider == "anthropic"

    def test_cross_region_inference_profile_unaffected(self, mock_boto3):
        llm = _make_llm("us.anthropic.claude-haiku-4-5-20251001-v1:0", mock_boto3)
        assert llm.provider == "anthropic"

    def test_arn_with_misspelled_provider_override_raises(self, mock_boto3):
        with pytest.raises(ValueError, match="Unknown provider_override 'anthorpic'"):
            _make_llm(self.ARN, mock_boto3, provider_override="anthorpic")


# ---------------------------------------------------------------------------
# LlmFactory
# ---------------------------------------------------------------------------

class TestLlmFactory:
    def test_aws_bedrock_uses_aws_bedrock_config(self):
        _, config_class = LlmFactory.provider_to_class["aws_bedrock"]
        assert config_class is AWSBedrockConfig

    def test_factory_create_accepts_aws_region(self, mock_boto3):
        """LlmFactory.create must not crash when aws_region is in the config dict.

        Before the fix, the factory mapped aws_bedrock to BaseLlmConfig which has
        no aws_region parameter, causing: TypeError: __init__() got an unexpected
        keyword argument 'aws_region'.
        """
        user_dict = {
            "model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
            "aws_region": "us-east-2",
            "temperature": 0.1,
            "max_tokens": 2000,
        }
        mock_boto3.converse.return_value = _converse_response()

        # This must not raise TypeError
        llm = LlmFactory.create("aws_bedrock", user_dict)

        assert isinstance(llm.config, AWSBedrockConfig)
        assert llm.config.aws_region == "us-east-2"
        assert llm.config.top_p is None


# ---------------------------------------------------------------------------
# _build_inference_config
# ---------------------------------------------------------------------------

class TestBuildInferenceConfig:
    """
    Unit tests for the _build_inference_config helper.
    Validates the exact keys present in the returned dict.
    """

    def test_anthropic_only_temperature_by_default(self, mock_boto3):
        llm = _make_llm("anthropic.claude-3-5-sonnet-20240620-v1:0", mock_boto3, temperature=0.5)
        cfg = llm._build_inference_config()
        assert "temperature" in cfg
        assert cfg["temperature"] == 0.5
        assert "topP" not in cfg, "topP must be absent when top_p not configured"

    def test_anthropic_top_p_explicitly_set_still_omits_top_p(self, mock_boto3):
        # Anthropic rejects both; topP must still be omitted even when set
        llm = _make_llm(
            "anthropic.claude-3-5-sonnet-20240620-v1:0",
            mock_boto3,
            temperature=0.5,
            top_p=0.9,
        )
        cfg = llm._build_inference_config()
        assert "temperature" in cfg
        assert "topP" not in cfg

    def test_anthropic_inference_profile_omits_top_p(self, mock_boto3):
        # Cross-region inference profiles (us.anthropic.*) follow the same rule
        llm = _make_llm("us.anthropic.claude-haiku-4-5-20251001-v1:0", mock_boto3, temperature=0.1)
        cfg = llm._build_inference_config()
        assert "topP" not in cfg

    def test_amazon_includes_top_p_when_set(self, mock_boto3):
        llm = _make_llm(
            "amazon.nova-3-mini-20241119-v1:0",
            mock_boto3,
            temperature=0.5,
            top_p=0.85,
        )
        cfg = llm._build_inference_config()
        assert cfg["temperature"] == 0.5
        assert cfg["topP"] == 0.85

    def test_amazon_omits_top_p_when_not_set(self, mock_boto3):
        llm = _make_llm("amazon.nova-3-mini-20241119-v1:0", mock_boto3, temperature=0.5)
        cfg = llm._build_inference_config()
        assert "topP" not in cfg

    def test_max_tokens_present(self, mock_boto3):
        llm = _make_llm(
            "anthropic.claude-3-5-sonnet-20240620-v1:0",
            mock_boto3,
            max_tokens=1024,
        )
        cfg = llm._build_inference_config()
        assert cfg["maxTokens"] == 1024

    def test_nova_fallback_max_tokens_when_absent(self, mock_boto3):
        """Legacy Nova Converse used 5000 when max_tokens was missing from the dict."""
        llm = _make_llm("amazon.nova-3-mini-20241119-v1:0", mock_boto3)
        llm.model_config.pop("max_tokens", None)
        cfg = llm._build_inference_config()
        assert cfg["maxTokens"] == 5000

    def test_anthropic_fallback_max_tokens_when_absent(self, mock_boto3):
        llm = _make_llm("anthropic.claude-3-5-sonnet-20240620-v1:0", mock_boto3)
        llm.model_config.pop("max_tokens", None)
        cfg = llm._build_inference_config()
        assert cfg["maxTokens"] == 2000

    def test_minimax_omits_top_p_when_explicitly_set(self, mock_boto3):
        # MiniMax M2.x (reasoning model) rejects both temperature and topP simultaneously.
        # Even when the user explicitly configures top_p, it must be omitted.
        llm = _make_llm(
            "minimax.minimax-m2.5",
            mock_boto3,
            temperature=0.1,
            top_p=0.9,
        )
        cfg = llm._build_inference_config()
        assert "temperature" in cfg
        assert "topP" not in cfg, "topP must be absent for MiniMax reasoning models"

    def test_minimax_only_temperature_by_default(self, mock_boto3):
        llm = _make_llm("minimax.minimax-m2.5", mock_boto3, temperature=0.1)
        cfg = llm._build_inference_config()
        assert cfg["temperature"] == 0.1
        assert "topP" not in cfg


# ---------------------------------------------------------------------------
# generate_response — Converse API call assertions
# ---------------------------------------------------------------------------

MESSAGES = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Hello"},
]

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "add_memory",
            "description": "Store a memory",
            "parameters": {
                "type": "object",
                "properties": {"data": {"type": "string"}},
                "required": ["data"],
            },
        },
    }
]


class TestGenerateResponseConverse:
    def test_standard_anthropic_no_top_p_in_converse_call(self, mock_boto3):
        mock_boto3.converse.return_value = _converse_response()
        llm = _make_llm("anthropic.claude-3-5-sonnet-20240620-v1:0", mock_boto3, temperature=0.2)

        llm.generate_response(MESSAGES)

        _, kwargs = mock_boto3.converse.call_args
        inference_cfg = kwargs["inferenceConfig"]
        assert "topP" not in inference_cfg
        assert inference_cfg["temperature"] == 0.2

    def test_standard_anthropic_inference_profile_no_top_p(self, mock_boto3):
        mock_boto3.converse.return_value = _converse_response()
        llm = _make_llm("us.anthropic.claude-haiku-4-5-20251001-v1:0", mock_boto3, temperature=0.1)

        llm.generate_response(MESSAGES)

        _, kwargs = mock_boto3.converse.call_args
        assert "topP" not in kwargs["inferenceConfig"]

    def test_with_tools_anthropic_no_top_p(self, mock_boto3):
        mock_boto3.converse.return_value = _converse_response()
        llm = _make_llm("anthropic.claude-3-5-sonnet-20240620-v1:0", mock_boto3, temperature=0.2)

        llm.generate_response(MESSAGES, tools=TOOLS)

        _, kwargs = mock_boto3.converse.call_args
        assert "topP" not in kwargs["inferenceConfig"]

    def test_with_tools_anthropic_top_p_set_still_omitted(self, mock_boto3):
        # Even if user explicitly sets top_p, Anthropic must not receive topP
        mock_boto3.converse.return_value = _converse_response()
        llm = _make_llm(
            "anthropic.claude-3-5-sonnet-20240620-v1:0",
            mock_boto3,
            temperature=0.2,
            top_p=0.9,
        )

        llm.generate_response(MESSAGES, tools=TOOLS)

        _, kwargs = mock_boto3.converse.call_args
        assert "topP" not in kwargs["inferenceConfig"]

    def test_nova_includes_top_p_when_explicitly_set(self, mock_boto3):
        mock_boto3.converse.return_value = _converse_response()
        llm = _make_llm(
            "amazon.nova-3-mini-20241119-v1:0",
            mock_boto3,
            temperature=0.5,
            top_p=0.85,
        )

        llm.generate_response(MESSAGES)

        _, kwargs = mock_boto3.converse.call_args
        assert kwargs["inferenceConfig"]["topP"] == 0.85

    def test_nova_omits_top_p_when_not_set(self, mock_boto3):
        mock_boto3.converse.return_value = _converse_response()
        llm = _make_llm("amazon.nova-3-mini-20241119-v1:0", mock_boto3, temperature=0.5)

        llm.generate_response(MESSAGES)

        _, kwargs = mock_boto3.converse.call_args
        assert "topP" not in kwargs["inferenceConfig"]

    def test_anthropic_model_kwargs_top_p_still_omits_top_p_in_converse(self, mock_boto3):
        """top_p injected via model_kwargs must not add topP for Anthropic Converse."""
        mock_boto3.converse.return_value = _converse_response()
        llm = _make_llm(
            "anthropic.claude-3-5-sonnet-20240620-v1:0",
            mock_boto3,
            temperature=0.2,
            model_kwargs={"top_p": 0.88},
        )

        llm.generate_response(MESSAGES)

        _, kwargs = mock_boto3.converse.call_args
        assert "topP" not in kwargs["inferenceConfig"]


# ---------------------------------------------------------------------------
# MiniMax provider
# ---------------------------------------------------------------------------

class TestMiniMaxProvider:
    """Tests for MiniMax models via Bedrock Converse API."""

    def test_extract_provider(self):
        assert extract_provider("minimax.minimax-m2.5") == "minimax"
        assert extract_provider("minimax.minimax-m2") == "minimax"

    def test_generate_response_text_only(self, mock_boto3):
        """Standard response: single text block."""
        mock_boto3.converse.return_value = _converse_response("Hello!")
        llm = _make_llm("minimax.minimax-m2.5", mock_boto3)

        result = llm.generate_response([{"role": "user", "content": "say hi"}])

        assert result == "Hello!"
        _, kwargs = mock_boto3.converse.call_args
        assert kwargs["modelId"] == "minimax.minimax-m2.5"
        assert kwargs["messages"][0]["role"] == "user"
        assert kwargs["messages"][0]["content"][0]["text"] == "say hi"

    def test_generate_response_reasoning_model(self, mock_boto3):
        """MiniMax M2.5 is a reasoning model: reasoningContent block comes before text."""
        reasoning_response = {
            "output": {
                "message": {
                    "content": [
                        {"reasoningContent": {"reasoningText": {"text": "Let me think..."}}},
                        {"text": "Hello!"},
                    ]
                }
            }
        }
        mock_boto3.converse.return_value = reasoning_response
        llm = _make_llm("minimax.minimax-m2.5", mock_boto3)

        result = llm.generate_response([{"role": "user", "content": "say hi"}])

        # Must skip reasoningContent and return the actual text block
        assert result == "Hello!"

    def test_inference_config(self, mock_boto3):
        """inferenceConfig should include maxTokens and temperature; no topP."""
        mock_boto3.converse.return_value = _converse_response()
        llm = _make_llm("minimax.minimax-m2.5", mock_boto3, temperature=0.2, max_tokens=512)

        llm.generate_response([{"role": "user", "content": "hi"}])

        _, kwargs = mock_boto3.converse.call_args
        assert kwargs["inferenceConfig"]["maxTokens"] == 512
        assert kwargs["inferenceConfig"]["temperature"] == 0.2
        assert "topP" not in kwargs["inferenceConfig"]

    def test_system_prompt_passed_correctly(self, mock_boto3):
        """System messages must be sent via top-level `system` param, not as a message role."""
        mock_boto3.converse.return_value = _converse_response('{"facts": ["test"]}')
        llm = _make_llm("minimax.minimax-m2.5", mock_boto3)

        llm.generate_response([
            {"role": "system", "content": "Return JSON only."},
            {"role": "user", "content": "Extract facts from: test"},
        ])

        _, kwargs = mock_boto3.converse.call_args
        # system prompt must be in top-level "system" key
        assert "system" in kwargs
        assert kwargs["system"][0]["text"] == "Return JSON only."
        # messages list must NOT contain a system role entry
        for msg in kwargs["messages"]:
            assert msg["role"] != "system"
        # user message must be present
        assert kwargs["messages"][0]["role"] == "user"


# ---------------------------------------------------------------------------
# _parse_response — legacy InvokeModel provider-specific parsing
# ---------------------------------------------------------------------------

class TestParseResponseLegacy:
    def test_ai21_missing_completions_returns_empty(self, mock_boto3):
        """When AI21 response lacks 'completions', the fallback default must
        be a valid dict (not a set literal), returning empty string."""
        llm = _make_llm("ai21.j2-mid-v1", mock_boto3)
        import io
        import json
        body = io.BytesIO(json.dumps({"not_completions": True}).encode())
        response = {"body": body}
        result = llm._parse_response(response, tools=None)
        assert result == ""

    def test_ai21_normal_response(self, mock_boto3):
        llm = _make_llm("ai21.j2-mid-v1", mock_boto3)
        import io
        import json
        body = io.BytesIO(json.dumps({
            "completions": [{"data": {"text": "hello from ai21"}}]
        }).encode())
        response = {"body": body}
        result = llm._parse_response(response, tools=None)
        assert result == "hello from ai21"


class TestAnthropicConverseContentParsing:
    """The Anthropic Converse branch must not assume content[0] is the text
    block: Claude reasoning models emit a reasoningContent block before the
    text block, and some stop conditions produce an empty content array. The
    parser iterates for the first block carrying text, like the MiniMax branch.
    """

    def test_text_after_reasoning_content_block(self, mock_boto3):
        mock_boto3.converse.return_value = {
            "output": {
                "message": {
                    "content": [
                        {"reasoningContent": {"reasoningText": {"text": "step by step..."}}},
                        {"text": "final answer"},
                    ]
                }
            }
        }
        llm = _make_llm("anthropic.claude-3-5-sonnet-20240620-v1:0", mock_boto3)

        assert llm.generate_response(MESSAGES) == "final answer"

    def test_empty_content_returns_empty_string(self, mock_boto3):
        mock_boto3.converse.return_value = {"output": {"message": {"content": []}}}
        llm = _make_llm("anthropic.claude-3-5-sonnet-20240620-v1:0", mock_boto3)

        assert llm.generate_response(MESSAGES) == ""

    def test_plain_text_content_still_returned(self, mock_boto3):
        mock_boto3.converse.return_value = _converse_response("plain answer")
        llm = _make_llm("anthropic.claude-3-5-sonnet-20240620-v1:0", mock_boto3)

        assert llm.generate_response(MESSAGES) == "plain answer"

    def test_object_style_response_iterates_blocks(self, mock_boto3):
        # Defensive attr-style branch: object wrapper with a reasoning block first.
        from types import SimpleNamespace

        reasoning_block = SimpleNamespace(reasoningContent={"reasoningText": {"text": "hmm"}})
        text_block = SimpleNamespace(text="object answer")
        mock_boto3.converse.return_value = SimpleNamespace(
            output=SimpleNamespace(message=SimpleNamespace(content=[reasoning_block, text_block]))
        )
        llm = _make_llm("anthropic.claude-3-5-sonnet-20240620-v1:0", mock_boto3)

        assert llm.generate_response(MESSAGES) == "object answer"