from __future__ import annotations

import json
import sys
from pathlib import Path

import krim.mcp as mcp_module
from krim.mcp import McpServer, McpServerConfig, load_mcp_config


def test_mcp_server_falls_back_to_line_stdio(tmp_path: Path, monkeypatch) -> None:
    monkeypatch.setattr(mcp_module, "MCP_READ_TIMEOUT", 2)

    server_script = tmp_path / "fake_line_mcp.py"
    server_script.write_text(
        """
import json
import sys

for raw in sys.stdin:
    line = raw.strip()
    if not line:
        continue
    try:
        msg = json.loads(line)
    except Exception:
        sys.stdout.write(json.dumps({
            "jsonrpc": "2.0",
            "id": None,
            "error": {"code": -32700, "message": "Parse error"},
        }) + "\\n")
        sys.stdout.flush()
        continue

    method = msg.get("method")
    if method == "initialize":
        sys.stdout.write(json.dumps({
            "jsonrpc": "2.0",
            "id": msg["id"],
            "result": {"protocolVersion": "2024-11-05", "capabilities": {}},
        }) + "\\n")
    elif method == "notifications/initialized":
        continue
    elif method == "tools/list":
        sys.stdout.write(json.dumps({
            "jsonrpc": "2.0",
            "id": msg["id"],
            "result": {
                "tools": [{
                    "name": "mcp__time__get_current_time",
                    "description": "Return a fixed time",
                    "inputSchema": {"type": "object", "properties": {}, "required": []},
                }]
            },
        }) + "\\n")
    elif method == "tools/call":
        sys.stdout.write(json.dumps({
            "jsonrpc": "2.0",
            "id": msg["id"],
            "result": {
                "content": [{"type": "text", "text": "2026-03-16T00:00:00Z"}],
            },
        }) + "\\n")
    else:
        sys.stdout.write(json.dumps({
            "jsonrpc": "2.0",
            "id": msg.get("id"),
            "result": {},
        }) + "\\n")
    sys.stdout.flush()
        """,
        encoding="utf-8",
    )

    server = McpServer(
        McpServerConfig(
            name="time",
            command=[sys.executable, str(server_script)],
        )
    )

    try:
        server.start()
        assert server.tools
        assert server.tools[0].name == "mcp__time__get_current_time"
        result = server.call_tool("mcp__time__get_current_time", {})
        assert "2026-03-16T00:00:00Z" in result
    finally:
        server.stop()


def test_mcp_server_line_mode_ignores_stdout_banner(tmp_path: Path, monkeypatch) -> None:
    monkeypatch.setattr(mcp_module, "MCP_READ_TIMEOUT", 2)

    server_script = tmp_path / "fake_banner_line_mcp.py"
    server_script.write_text(
        """
import json
import sys

print("Starting banner on stdout...")
sys.stdout.flush()

for raw in sys.stdin:
    line = raw.strip()
    if not line:
        continue

    msg = json.loads(line)
    method = msg.get("method")
    if method == "initialize":
        sys.stdout.write(json.dumps({
            "jsonrpc": "2.0",
            "id": msg["id"],
            "result": {"protocolVersion": "2024-11-05", "capabilities": {}},
        }) + "\\n")
    elif method == "notifications/initialized":
        continue
    elif method == "tools/list":
        sys.stdout.write(json.dumps({
            "jsonrpc": "2.0",
            "id": msg["id"],
            "result": {
                "tools": [{
                    "name": "mcp__banner__ping",
                    "description": "Ping tool",
                    "inputSchema": {"type": "object", "properties": {}, "required": []},
                }]
            },
        }) + "\\n")
    elif method == "tools/call":
        sys.stdout.write(json.dumps({
            "jsonrpc": "2.0",
            "id": msg["id"],
            "result": {"content": [{"type": "text", "text": "pong"}]},
        }) + "\\n")
    sys.stdout.flush()
        """,
        encoding="utf-8",
    )

    server = McpServer(
        McpServerConfig(
            name="banner",
            command=[sys.executable, str(server_script)],
            wire_format="line",
        )
    )

    try:
        server.start()
        assert [tool.name for tool in server.tools] == ["mcp__banner__ping"]
        assert server.call_tool("mcp__banner__ping", {}) == "pong"
    finally:
        server.stop()


def test_load_mcp_config_supports_command_args_and_wire_format(tmp_path: Path) -> None:
    config_dir = tmp_path / ".krim"
    config_dir.mkdir()
    (config_dir / "mcp.json").write_text(
        json.dumps(
            {
                "mcpServers": {
                    "context7": {
                        "command": "npx",
                        "args": ["-y", "@upstash/context7-mcp"],
                        "wire_format": "line",
                        "env": {"NODE_ENV": "production"},
                    },
                    "custom": {
                        "command": ["python", "-m", "custom_server"],
                        "args": ["--port", "8123"],
                    },
                }
            }
        ),
        encoding="utf-8",
    )

    configs = load_mcp_config(global_dir=config_dir)

    assert [(config.name, config.command, config.wire_format) for config in configs] == [
        ("context7", ["npx", "-y", "@upstash/context7-mcp"], "line"),
        ("custom", ["python", "-m", "custom_server", "--port", "8123"], "auto"),
    ]
    assert configs[0].env == {"NODE_ENV": "production"}


def test_mcp_server_aliases_invalid_tool_names(tmp_path: Path, monkeypatch) -> None:
    monkeypatch.setattr(mcp_module, "MCP_READ_TIMEOUT", 2)

    server_script = tmp_path / "fake_invalid_name_mcp.py"
    server_script.write_text(
        """
import json
import sys

for raw in sys.stdin:
    line = raw.strip()
    if not line:
        continue
    msg = json.loads(line)
    method = msg.get("method")
    if method == "initialize":
        sys.stdout.write(json.dumps({
            "jsonrpc": "2.0",
            "id": msg["id"],
            "result": {"protocolVersion": "2024-11-05", "capabilities": {}},
        }) + "\\n")
    elif method == "notifications/initialized":
        continue
    elif method == "tools/list":
        sys.stdout.write(json.dumps({
            "jsonrpc": "2.0",
            "id": msg["id"],
            "result": {
                "tools": [{
                    "name": "KRIS API MCP",
                    "description": "Remote KRIS tool",
                    "inputSchema": {"type": "object", "properties": {}, "required": []},
                }]
            },
        }) + "\\n")
    elif method == "tools/call":
        sys.stdout.write(json.dumps({
            "jsonrpc": "2.0",
            "id": msg["id"],
            "result": {"content": [{"type": "text", "text": msg["params"]["name"]}]},
        }) + "\\n")
    sys.stdout.flush()
        """,
        encoding="utf-8",
    )

    server = McpServer(
        McpServerConfig(
            name="kris",
            command=[sys.executable, str(server_script)],
            wire_format="line",
        )
    )

    try:
        server.start()
        assert [tool.name for tool in server.tools] == ["kris_api_mcp"]
        assert "original MCP tool name: KRIS API MCP" in server.tools[0].description
        assert server.tools[0].run() == "KRIS API MCP"
    finally:
        server.stop()
