"""Tests for the persistent MCP manager."""

import json
from pathlib import Path

from pydantic import BaseModel

from yier_agents.src.mcp_manager import MCPManager
from yier_agents.src.tool import Tool, ToolContext, ToolOutput


class EmptyParams(BaseModel):
    """Simple params model for fake MCP tools."""


def make_tool(name: str) -> Tool:
    """Create one fake tool with a stable name."""

    async def execute(params: EmptyParams, ctx: ToolContext) -> ToolOutput:
        return ToolOutput(content=f"{name} ok")

    return Tool(
        name=name,
        description=f"{name} description",
        parameters=EmptyParams,
        execute=execute,
    )


class FakeClient:
    """Test double that behaves like a connected MCP client."""

    def __init__(self, tool_name: str) -> None:
        self.tool_name = tool_name
        self.closed = False
        self.enter_count = 0

    async def __aenter__(self):
        self.enter_count += 1
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        self.closed = True

    async def get_tools(self) -> list[Tool]:
        return [make_tool(self.tool_name)]


def write_mcp_config(config_dir: Path, payload: dict) -> None:
    """Persist test MCP config to the managed config file."""
    config_dir.mkdir(parents=True, exist_ok=True)
    config_path = config_dir / ".yier.json"
    config_path.write_text(json.dumps({"mcpServers": payload}), encoding="utf-8")


def make_client_factory(created: dict[str, FakeClient]):
    """Create fake MCP clients keyed by a readable server signature."""

    def factory(server_cfg):
        if server_cfg["type"] == "stdio":
            label = f"{server_cfg['command']}:{','.join(server_cfg.get('args', []))}"
        else:
            label = server_cfg["url"]
        client = FakeClient(label.replace("/", "_"))
        created[label] = client
        return client

    return factory


async def test_manager_connects_servers_and_returns_tools(tmp_path: Path):
    config_dir = tmp_path / "config"
    write_mcp_config(
        config_dir,
        {
            "docs": {
                "type": "stdio",
                "command": "uvx",
                "args": ["docs-mcp"],
            }
        },
    )
    created: dict[str, FakeClient] = {}
    manager = MCPManager(
        config_dir=config_dir,
        client_factory=make_client_factory(created),
    )

    await manager.reload(force_reconnect=True)

    tools = await manager.get_tools()
    status = await manager.get_status()

    assert [tool.name for tool in tools] == ["uvx:docs-mcp"]
    assert status["docs"]["status"] == "connected"
    assert status["docs"]["tool_count"] == 1
    assert "uvx:docs-mcp" in created

    await manager.stop()
    assert created["uvx:docs-mcp"].closed is True


async def test_manager_reload_reconnects_changed_and_removed_servers(tmp_path: Path):
    config_dir = tmp_path / "config"
    write_mcp_config(
        config_dir,
        {
            "alpha": {
                "type": "stdio",
                "command": "uvx",
                "args": ["alpha-v1"],
            },
            "beta": {
                "type": "stdio",
                "command": "uvx",
                "args": ["beta"],
            },
        },
    )
    created: dict[str, FakeClient] = {}
    manager = MCPManager(
        config_dir=config_dir,
        client_factory=make_client_factory(created),
    )

    await manager.reload(force_reconnect=True)
    alpha_v1 = created["uvx:alpha-v1"]
    beta_v1 = created["uvx:beta"]

    write_mcp_config(
        config_dir,
        {
            "alpha": {
                "type": "stdio",
                "command": "uvx",
                "args": ["alpha-v2"],
            },
            "gamma": {
                "type": "stdio",
                "command": "uvx",
                "args": ["gamma"],
            },
        },
    )

    changed = await manager.reload_if_changed()
    tools = sorted(tool.name for tool in await manager.get_tools())
    status = await manager.get_status()

    assert changed is True
    assert alpha_v1.closed is True
    assert beta_v1.closed is True
    assert tools == ["uvx:alpha-v2", "uvx:gamma"]
    assert status["alpha"]["status"] == "connected"
    assert status["gamma"]["status"] == "connected"
    assert "beta" not in status

    await manager.stop()


async def test_manager_skips_disabled_servers(tmp_path: Path):
    config_dir = tmp_path / "config"
    write_mcp_config(
        config_dir,
        {
            "disabled-server": {
                "type": "stdio",
                "command": "uvx",
                "args": ["ignored"],
                "enabled": False,
            }
        },
    )
    created: dict[str, FakeClient] = {}
    manager = MCPManager(
        config_dir=config_dir,
        client_factory=make_client_factory(created),
    )

    await manager.reload(force_reconnect=True)

    assert await manager.get_tools() == []
    status = await manager.get_status()
    assert status["disabled-server"]["status"] == "disabled"
    assert created == {}
