import asyncio

import pytest

from mcpbridge.adapters._openai_style import append_tool_results_openai, extract_text_openai, extract_tool_calls_openai, is_done_openai
from mcpbridge.adapters.base import BaseLLMAdapter, ToolCall, ToolResult
from mcpbridge.config import LoopConfig, PromptConfig
from mcpbridge.loop import run_loop
from mcpbridge.prompt_builder import PromptBuilder
from mcpbridge.registry import DynamicToolRegistry, MCPTool
from mcpbridge.transport.base import BaseTransport


class ToolTransport(BaseTransport):
    """Fake MCP transport that only supports `tools/call`."""

    def __init__(self, server_name: str) -> None:
        self._server_name = server_name

    @property
    def is_connected(self) -> bool:
        return True

    async def connect(self) -> None:
        return

    async def close(self) -> None:
        return

    async def send_notification(self, method: str, params: dict) -> None:
        return

    async def send(self, message: dict) -> dict:
        if message.get("method") != "tools/call":
            raise AssertionError("Unexpected method")
        params = message.get("params") or {}
        name = params.get("name")
        args = params.get("arguments") or {}
        if name == "echo":
            text = args.get("text", "")
            return {
                "jsonrpc": "2.0",
                "id": message.get("id"),
                "result": {"content": [{"type": "text", "text": f"ECHO:{text}"}]},
            }
        if name == "upper":
            text = args.get("text", "")
            return {
                "jsonrpc": "2.0",
                "id": message.get("id"),
                "result": {"content": [{"type": "text", "text": f"UPPER:{text.upper()}"}]},
            }
        return {"jsonrpc": "2.0", "id": message.get("id"), "result": {"content": []}}


class DummyOpenAIStyleAdapter(BaseLLMAdapter):
    """Adapter that returns a tool call once, then returns final text."""

    def __init__(self) -> None:
        self._provider = "openai_compatible"
        self._call_count = 0

    @property
    def provider_name(self) -> str:
        return self._provider

    @property
    def supports_parallel_tool_calls(self) -> bool:
        return True

    @property
    def supports_system_prompt(self) -> bool:
        return False

    async def chat(self, messages, tools, system=None, **kwargs):
        self._call_count += 1
        if self._call_count == 1:
            return {
                "choices": [
                    {
                        "message": {
                            "content": None,
                            "tool_calls": [
                                {
                                    "id": "call1",
                                    "function": {"name": "srv__echo", "arguments": '{"text":"hi"}'},
                                }
                            ],
                        }
                    }
                ]
            }
        return {"choices": [{"message": {"content": "done", "tool_calls": []}}]}

    def extract_tool_calls(self, response):
        return extract_tool_calls_openai(response)

    def extract_text(self, response):
        return extract_text_openai(response)

    def is_done(self, response):
        return is_done_openai(response)

    def append_tool_results(self, messages, response, tool_results):
        return append_tool_results_openai(messages, response, tool_results)

    def count_tokens(self, messages, system):
        return 0


def test_run_loop_single_tool_happy_path():
    async def _run():
        adapter = DummyOpenAIStyleAdapter()

        # Registry with a single tool.
        reg = DynamicToolRegistry(server_name="merged", namespace_strategy="prefix")
        reg.tools = {
            "srv__echo": MCPTool(
                name="echo",
                description="Echo",
                input_schema={
                    "type": "object",
                    "properties": {"text": {"type": "string"}},
                    "required": ["text"],
                },
                server_name="srv",
                namespaced_name="srv__echo",
            )
        }

        transport_map = {"srv": ToolTransport("srv")}

        pb = PromptBuilder(PromptConfig(system="SYS", interpolate=False))
        pb.set_system_in_messages(True)

        history: list[dict] = []
        loop_cfg = LoopConfig(max_iterations=3, max_tokens_total=32000, tool_timeout=2, parallel_tool_calls=True)
        res = await run_loop(
            adapter=adapter,
            prompt_builder=pb,
            registry=reg,
            transport=transport_map,
            user_query="Q",
            history=history,
            loop_config=loop_cfg,
        )
        assert res.finish_reason == "done"
        assert res.text == "done"
        assert len(res.tool_calls_made) == 1
        assert "ECHO:hi" in res.tool_calls_made[0]["result"]
        # History should be mutated by the loop.
        assert any(m.get("role") == "tool" for m in history)

    asyncio.run(_run())


class DummyParallelAdapter(DummyOpenAIStyleAdapter):
    async def chat(self, messages, tools, system=None, **kwargs):
        self._call_count += 1
        if self._call_count == 1:
            return {
                "choices": [
                    {
                        "message": {
                            "content": None,
                            "tool_calls": [
                                {
                                    "id": "call1",
                                    "function": {"name": "srv__echo", "arguments": '{"text":"hi"}'},
                                },
                                {
                                    "id": "call2",
                                    "function": {"name": "srv__upper", "arguments": '{"text":"wow"}'},
                                },
                            ],
                        }
                    }
                ]
            }
        return {"choices": [{"message": {"content": "done2", "tool_calls": []}}]}


def test_run_loop_parallel_tool_calls():
    async def _run():
        adapter = DummyParallelAdapter()
        reg = DynamicToolRegistry(server_name="merged", namespace_strategy="prefix")
        reg.tools = {
            "srv__echo": MCPTool(
                name="echo",
                description="Echo",
                input_schema={
                    "type": "object",
                    "properties": {"text": {"type": "string"}},
                    "required": ["text"],
                },
                server_name="srv",
                namespaced_name="srv__echo",
            ),
            "srv__upper": MCPTool(
                name="upper",
                description="Upper",
                input_schema={
                    "type": "object",
                    "properties": {"text": {"type": "string"}},
                    "required": ["text"],
                },
                server_name="srv",
                namespaced_name="srv__upper",
            ),
        }

        transport_map = {"srv": ToolTransport("srv")}
        pb = PromptBuilder(PromptConfig(system="SYS", interpolate=False))
        pb.set_system_in_messages(True)

        history: list[dict] = []
        loop_cfg = LoopConfig(max_iterations=3, max_tokens_total=32000, tool_timeout=2, parallel_tool_calls=True)
        res = await run_loop(
            adapter=adapter,
            prompt_builder=pb,
            registry=reg,
            transport=transport_map,
            user_query="Q",
            history=history,
            loop_config=loop_cfg,
        )
        assert res.finish_reason == "done"
        assert res.text == "done2"
        assert len(res.tool_calls_made) == 2
        assert any("UPPER:WOW" in x["result"] for x in res.tool_calls_made)

    asyncio.run(_run())


class DummyNeverDoneAdapter(BaseLLMAdapter):
    """Adapter that always returns a tool call (never reaches `is_done`)."""

    def __init__(self) -> None:
        self._provider = "openai_compatible"

    @property
    def provider_name(self) -> str:
        return self._provider

    @property
    def supports_parallel_tool_calls(self) -> bool:
        return True

    @property
    def supports_system_prompt(self) -> bool:
        return False

    async def chat(self, messages, tools, system=None, **kwargs):
        # Always request the same tool call; this simulates a model that
        # gets stuck re-calling tools.
        return {
            "choices": [
                {
                    "message": {
                        "content": None,
                        "tool_calls": [
                            {
                                "id": "call1",
                                "function": {"name": "srv__echo", "arguments": '{"text":"hi"}'},
                            }
                        ],
                    }
                }
            ]
        }

    def extract_tool_calls(self, response):
        return extract_tool_calls_openai(response)

    def extract_text(self, response):
        return ""

    def is_done(self, response):
        return False

    def append_tool_results(self, messages, response, tool_results):
        return append_tool_results_openai(messages, response, tool_results)

    def count_tokens(self, messages, system):
        return 0


def test_run_loop_max_iterations_returns_result_not_exception():
    async def _run():
        adapter = DummyNeverDoneAdapter()

        # Registry with a single tool.
        reg = DynamicToolRegistry(server_name="merged", namespace_strategy="prefix")
        reg.tools = {
            "srv__echo": MCPTool(
                name="echo",
                description="Echo",
                input_schema={
                    "type": "object",
                    "properties": {"text": {"type": "string"}},
                    "required": ["text"],
                },
                server_name="srv",
                namespaced_name="srv__echo",
            )
        }

        transport_map = {"srv": ToolTransport("srv")}
        pb = PromptBuilder(PromptConfig(system="SYS", interpolate=False))
        pb.set_system_in_messages(True)

        history: list[dict] = []
        loop_cfg = LoopConfig(max_iterations=1, max_tokens_total=32000, tool_timeout=2, parallel_tool_calls=True)

        res = await run_loop(
            adapter=adapter,
            prompt_builder=pb,
            registry=reg,
            transport=transport_map,
            user_query="Q",
            history=history,
            loop_config=loop_cfg,
        )

        assert res.finish_reason == "max_iterations"
        assert res.iterations == 1
        assert len(res.tool_calls_made) == 1
        assert "ECHO:hi" in (res.text or "")

    asyncio.run(_run())

