import base64
from collections.abc import Iterator
from decimal import Decimal
from typing import Any
from unittest.mock import Mock, patch

import pytest
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session

from core.mcp.types import (
    AudioContent,
    BlobResourceContents,
    CallToolResult,
    EmbeddedResource,
    ImageContent,
    TextContent,
    TextResourceContents,
)
from core.tools.__base.tool_runtime import ToolRuntime
from core.tools.entities.common_entities import I18nObject
from core.tools.entities.tool_entities import ToolEntity, ToolIdentity, ToolInvokeMessage
from core.tools.mcp_tool.tool import MCPTool
from graphon.model_runtime.entities.llm_entities import LLMUsage


@pytest.fixture
def orm_session(sqlite_engine: Engine) -> Iterator[Session]:
    """Use a real ORM session while MCP transport remains mocked."""
    with Session(sqlite_engine) as session:
        yield session


def _make_mcp_tool(output_schema: dict[str, Any] | None = None) -> MCPTool:
    identity = ToolIdentity(
        author="test",
        name="test_mcp_tool",
        label=I18nObject(en_US="Test MCP Tool", zh_Hans="测试MCP工具"),
        provider="test_provider",
    )
    entity = ToolEntity(identity=identity, output_schema=output_schema or {})
    runtime = Mock(spec=ToolRuntime)
    runtime.credentials = {}
    return MCPTool(
        entity=entity,
        runtime=runtime,
        tenant_id="test_tenant",
        icon="",
        server_url="https://server.invalid",
        provider_id="provider_1",
        headers={},
    )


class TestMCPToolInvoke:
    @pytest.mark.parametrize(
        ("content_factory", "mime_type"),
        [
            (
                lambda b64, mt: ImageContent(type="image", data=b64, mimeType=mt),
                "image/png",
            ),
            (
                lambda b64, mt: AudioContent(type="audio", data=b64, mimeType=mt),
                "audio/mpeg",
            ),
        ],
    )
    def test_invoke_image_or_audio_yields_blob(self, content_factory, mime_type, orm_session: Session) -> None:
        tool = _make_mcp_tool()
        raw = b"\x00\x01test-bytes\x02"
        b64 = base64.b64encode(raw).decode()
        content = content_factory(b64, mime_type)
        result = CallToolResult(content=[content])

        with patch.object(tool, "invoke_remote_mcp_tool", return_value=result):
            messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={}))

        assert len(messages) == 1
        msg = messages[0]
        assert msg.type == ToolInvokeMessage.MessageType.BLOB
        assert isinstance(msg.message, ToolInvokeMessage.BlobMessage)
        assert msg.message.blob == raw
        assert msg.meta == {"mime_type": mime_type}

    def test_invoke_embedded_text_resource_yields_text(self, orm_session: Session) -> None:
        tool = _make_mcp_tool()
        text_resource = TextResourceContents(uri="file://test.txt", mimeType="text/plain", text="hello world")
        content = EmbeddedResource(type="resource", resource=text_resource)
        result = CallToolResult(content=[content])

        with patch.object(tool, "invoke_remote_mcp_tool", return_value=result):
            messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={}))

        assert len(messages) == 1
        msg = messages[0]
        assert msg.type == ToolInvokeMessage.MessageType.TEXT
        assert isinstance(msg.message, ToolInvokeMessage.TextMessage)
        assert msg.message.text == "hello world"

    @pytest.mark.parametrize(
        ("mime_type", "expected_mime"),
        [("application/pdf", "application/pdf"), (None, "application/octet-stream")],
    )
    def test_invoke_embedded_blob_resource_yields_blob(self, mime_type, expected_mime, orm_session: Session) -> None:
        tool = _make_mcp_tool()
        raw = b"binary-data"
        b64 = base64.b64encode(raw).decode()
        blob_resource = BlobResourceContents(uri="file://doc.bin", mimeType=mime_type, blob=b64)
        content = EmbeddedResource(type="resource", resource=blob_resource)
        result = CallToolResult(content=[content])

        with patch.object(tool, "invoke_remote_mcp_tool", return_value=result):
            messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={}))

        assert len(messages) == 1
        msg = messages[0]
        assert msg.type == ToolInvokeMessage.MessageType.BLOB
        assert isinstance(msg.message, ToolInvokeMessage.BlobMessage)
        assert msg.message.blob == raw
        assert msg.meta == {"mime_type": expected_mime}

    def test_invoke_yields_variables_when_structured_content_and_schema(self, orm_session: Session) -> None:
        tool = _make_mcp_tool(output_schema={"type": "object"})
        result = CallToolResult(content=[], structuredContent={"a": 1, "b": "x"})

        with patch.object(tool, "invoke_remote_mcp_tool", return_value=result):
            messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={}))

        # Expect two variable messages corresponding to keys a and b
        assert len(messages) == 2
        var_msgs = [m for m in messages if isinstance(m.message, ToolInvokeMessage.VariableMessage)]
        assert {m.message.variable_name for m in var_msgs} == {"a", "b"}
        # Validate values
        values = {m.message.variable_name: m.message.variable_value for m in var_msgs}
        assert values == {"a": 1, "b": "x"}


class TestMCPToolUsageExtraction:
    """Test usage metadata extraction from MCP tool results."""

    def test_extract_usage_dict_from_direct_usage_field(self) -> None:
        """Test extraction when usage is directly in meta.usage field."""
        meta = {
            "usage": {
                "prompt_tokens": 100,
                "completion_tokens": 50,
                "total_tokens": 150,
                "total_price": "0.001",
                "currency": "USD",
            }
        }
        usage_dict = MCPTool._extract_usage_dict(meta)
        assert usage_dict is not None
        assert usage_dict["prompt_tokens"] == 100
        assert usage_dict["completion_tokens"] == 50
        assert usage_dict["total_tokens"] == 150
        assert usage_dict["total_price"] == "0.001"
        assert usage_dict["currency"] == "USD"

    def test_extract_usage_dict_from_nested_metadata(self) -> None:
        """Test extraction when usage is nested in meta.metadata.usage."""
        meta = {
            "metadata": {
                "usage": {
                    "prompt_tokens": 200,
                    "completion_tokens": 100,
                    "total_tokens": 300,
                }
            }
        }
        usage_dict = MCPTool._extract_usage_dict(meta)
        assert usage_dict is not None
        assert usage_dict["prompt_tokens"] == 200
        assert usage_dict["total_tokens"] == 300

    def test_extract_usage_dict_from_flat_token_fields(self) -> None:
        """Test extraction when token counts are directly in meta."""
        meta = {
            "prompt_tokens": 150,
            "completion_tokens": 75,
            "total_tokens": 225,
            "currency": "EUR",
        }
        usage_dict = MCPTool._extract_usage_dict(meta)
        assert usage_dict is not None
        assert usage_dict["prompt_tokens"] == 150
        assert usage_dict["completion_tokens"] == 75
        assert usage_dict["total_tokens"] == 225
        assert usage_dict["currency"] == "EUR"

    def test_extract_usage_dict_recursive(self) -> None:
        """Test recursive search through nested structures."""
        meta = {
            "custom": {
                "nested": {
                    "usage": {
                        "total_tokens": 500,
                        "prompt_tokens": 300,
                        "completion_tokens": 200,
                    }
                }
            }
        }
        usage_dict = MCPTool._extract_usage_dict(meta)
        assert usage_dict is not None
        assert usage_dict["total_tokens"] == 500

    def test_extract_usage_dict_from_list(self) -> None:
        """Test extraction from nested list structures."""
        meta = {
            "items": [
                {"usage": {"total_tokens": 100}},
                {"other": "data"},
            ]
        }
        usage_dict = MCPTool._extract_usage_dict(meta)
        assert usage_dict is not None
        assert usage_dict["total_tokens"] == 100

    def test_extract_usage_dict_returns_none_when_missing(self) -> None:
        """Test that None is returned when no usage data is present."""
        meta = {"other": "data", "custom": {"nested": {"value": 123}}}
        usage_dict = MCPTool._extract_usage_dict(meta)
        assert usage_dict is None

    def test_extract_usage_dict_empty_meta(self) -> None:
        """Test with empty meta dict."""
        usage_dict = MCPTool._extract_usage_dict({})
        assert usage_dict is None

    def test_derive_usage_from_result_with_meta(self) -> None:
        """Test _derive_usage_from_result with populated meta."""
        meta = {
            "usage": {
                "prompt_tokens": 100,
                "completion_tokens": 50,
                "total_tokens": 150,
                "total_price": "0.0015",
                "currency": "USD",
            }
        }
        result = CallToolResult(content=[], _meta=meta)
        usage = MCPTool._derive_usage_from_result(result)

        assert isinstance(usage, LLMUsage)
        assert usage.prompt_tokens == 100
        assert usage.completion_tokens == 50
        assert usage.total_tokens == 150
        assert usage.total_price == Decimal("0.0015")
        assert usage.currency == "USD"

    def test_derive_usage_from_result_without_meta(self) -> None:
        """Test _derive_usage_from_result with no meta returns empty usage."""
        result = CallToolResult(content=[], meta=None)
        usage = MCPTool._derive_usage_from_result(result)

        assert isinstance(usage, LLMUsage)
        assert usage.total_tokens == 0
        assert usage.prompt_tokens == 0
        assert usage.completion_tokens == 0

    def test_derive_usage_from_result_calculates_total_tokens(self) -> None:
        """Test that total_tokens is calculated when missing."""
        meta = {
            "usage": {
                "prompt_tokens": 100,
                "completion_tokens": 50,
                # total_tokens is missing
            }
        }
        result = CallToolResult(content=[], _meta=meta)
        usage = MCPTool._derive_usage_from_result(result)

        assert usage.total_tokens == 150  # 100 + 50
        assert usage.prompt_tokens == 100
        assert usage.completion_tokens == 50

    def test_invoke_sets_latest_usage_from_meta(self, orm_session: Session) -> None:
        """Test that _invoke sets _latest_usage from result meta."""
        tool = _make_mcp_tool()
        meta = {
            "usage": {
                "prompt_tokens": 200,
                "completion_tokens": 100,
                "total_tokens": 300,
                "total_price": "0.003",
                "currency": "USD",
            }
        }
        result = CallToolResult(content=[TextContent(type="text", text="test")], _meta=meta)

        with patch.object(tool, "invoke_remote_mcp_tool", return_value=result):
            list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={}))

        # Verify latest_usage was set correctly
        assert tool.latest_usage.prompt_tokens == 200
        assert tool.latest_usage.completion_tokens == 100
        assert tool.latest_usage.total_tokens == 300
        assert tool.latest_usage.total_price == Decimal("0.003")

    def test_invoke_with_no_meta_returns_empty_usage(self, orm_session: Session) -> None:
        """Test that _invoke returns empty usage when no meta is present."""
        tool = _make_mcp_tool()
        result = CallToolResult(content=[TextContent(type="text", text="test")], _meta=None)

        with patch.object(tool, "invoke_remote_mcp_tool", return_value=result):
            list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={}))

        # Verify latest_usage is empty
        assert tool.latest_usage.total_tokens == 0
        assert tool.latest_usage.prompt_tokens == 0
        assert tool.latest_usage.completion_tokens == 0

    def test_latest_usage_property_returns_llm_usage(self) -> None:
        """Test that latest_usage property returns LLMUsage instance."""
        tool = _make_mcp_tool()
        assert isinstance(tool.latest_usage, LLMUsage)

    def test_initial_usage_is_empty(self) -> None:
        """Test that MCPTool is initialized with empty usage."""
        tool = _make_mcp_tool()
        assert tool.latest_usage.total_tokens == 0
        assert tool.latest_usage.prompt_tokens == 0
        assert tool.latest_usage.completion_tokens == 0
        assert tool.latest_usage.total_price == Decimal(0)

    @pytest.mark.parametrize(
        "meta_data",
        [
            # Direct usage field
            {"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}},
            # Nested metadata
            {"metadata": {"usage": {"total_tokens": 100}}},
            # Flat token fields
            {"total_tokens": 50, "prompt_tokens": 30, "completion_tokens": 20},
            # With price info
            {
                "usage": {
                    "total_tokens": 150,
                    "total_price": "0.002",
                    "currency": "EUR",
                }
            },
            # Deep nested
            {"level1": {"level2": {"usage": {"total_tokens": 200}}}},
        ],
    )
    def test_various_meta_formats(self, meta_data) -> None:
        """Test that various meta formats are correctly parsed."""
        result = CallToolResult(content=[], _meta=meta_data)
        usage = MCPTool._derive_usage_from_result(result)

        assert isinstance(usage, LLMUsage)
        # Should have at least some usage data
        if meta_data.get("usage", {}).get("total_tokens") or meta_data.get("total_tokens"):
            expected_total = (
                meta_data.get("usage", {}).get("total_tokens")
                or meta_data.get("total_tokens")
                or meta_data.get("metadata", {}).get("usage", {}).get("total_tokens")
                or meta_data.get("level1", {}).get("level2", {}).get("usage", {}).get("total_tokens")
            )
            if expected_total:
                assert usage.total_tokens == expected_total