"""MCP 三模块覆盖补点（CI 第二阶段覆盖率）。

- websearch/imagegen：handle_request 全分支（协议方法、工具
  分发、未知工具、工具异常、非对象请求、通知）+ main() 的
  stdio 循环（坏 JSON、通知不回包），工具执行用注入的
  fake，不联网、不调真实模型/网关。
- playwright：全部 14 个工具分发 + 未知工具 + 工具异常 +
  协议方法 + _handle_line 的解析错误与 handler 异常分支。
  fake browser 只返回标记 dict，断言分发到了正确方法。
- flat-copy：三个模块 + 后端文件拷到隔离临时目录，在
  PYTHONPATH 只含该目录的子进程里 import——cloud 包不可见，
  必须走 flat-copy 回落；断言后端类的 __module__ 是顶层
  模块名（不是 cloud.*），即回落真实发生（importlib 重构
  的回归证据，见 F401/mypy no-redef 一节）。
"""

import io
import json
import os
import shutil
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path
from unittest import mock

try:
    from cloud.websearch_mcp import handle_request as ws_handle
    from cloud.websearch_mcp import main as ws_main
    from cloud.imagegen_mcp import handle_request as ig_handle
    from cloud.imagegen_mcp import main as ig_main
    from cloud import playwright_mcp as pw_mod
    from cloud.playwright_mcp import handle_request as pw_handle
    from cloud.playwright_mcp import _handle_line
except ImportError:
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from cloud.websearch_mcp import handle_request as ws_handle
    from cloud.websearch_mcp import main as ws_main
    from cloud.imagegen_mcp import handle_request as ig_handle
    from cloud.imagegen_mcp import main as ig_main
    from cloud import playwright_mcp as pw_mod
    from cloud.playwright_mcp import handle_request as pw_handle
    from cloud.playwright_mcp import _handle_line


def _call(handle, method, req_id=1, name=None, args=None):
    req = {"jsonrpc": "2.0", "id": req_id, "method": method}
    if name is not None:
        req["params"] = {"name": name, "arguments": args or {}}
    return handle(req)


class WebSearchMcpTest(unittest.TestCase):
    def test_protocol_methods(self):
        self.assertIsNone(ws_handle({"method": "notifications/initialized"}))
        self.assertIsNone(ws_handle({"method": "ping"}))
        ping = _call(ws_handle, "ping")
        self.assertEqual(ping["result"], {})
        init = {"jsonrpc": "2.0", "id": 1, "method": "initialize",
                "params": {"protocolVersion": "2025-01-01"}}
        resp = ws_handle(init)
        self.assertEqual(resp["result"]["protocolVersion"], "2025-01-01")
        self.assertEqual(resp["result"]["serverInfo"]["name"], "websearch")
        self.assertIsNone(ws_handle({"method": "initialize"}))
        listed = _call(ws_handle, "tools/list")
        self.assertIn("web_search",
                      [t["name"] for t in listed["result"]["tools"]])

    def test_tools_call_and_errors(self):
        fake_search = lambda q, limit=5, provider="auto", scope="general": {  # noqa: E731
            "results": [{"title": q}]}
        # 工具执行一律用注入的 fake，不调真实 search（不联网）。
        resp = ws_handle({"jsonrpc": "2.0", "id": 2, "method": "tools/call",
                          "params": {"name": "web_search",
                                     "arguments": {"query": "hello"}}},
                         search_fn=fake_search)
        self.assertFalse(resp["result"]["isError"])
        self.assertIn("hello", resp["result"]["content"][0]["text"])
        resp = ws_handle({"jsonrpc": "2.0", "id": 3, "method": "tools/call",
                          "params": {"name": "web_fetch",
                                     "arguments": {"url": "https://x/"}}},
                         fetch_fn=lambda u, char_limit=15000: {"text": "T"})
        self.assertIn("T", resp["result"]["content"][0]["text"])
        unknown = _call(ws_handle, "tools/call", name="nope")
        self.assertEqual(unknown["error"]["code"], -32602)
        boom = ws_handle({"jsonrpc": "2.0", "id": 4, "method": "tools/call",
                          "params": {"name": "web_search",
                                     "arguments": {"query": "q"}}},
                         search_fn=lambda *a, **k: (_ for _ in ()).throw(
                             RuntimeError("boom")))
        self.assertTrue(boom["result"]["isError"])
        self.assertIn("boom", boom["result"]["content"][0]["text"])
        self.assertEqual(ws_handle("nope")["error"]["code"], -32700)
        self.assertEqual(
            _call(ws_handle, "bogus/method")["error"]["code"], -32601)
        self.assertIsNone(ws_handle(
            {"method": "tools/call",
             "params": {"name": "web_search", "arguments": {}}}))

    def test_main_stdio_loop(self):
        lines = ('{"bad json\n'
                 '{"jsonrpc": "2.0", "id": 1, "method": "initialize",'
                 ' "params": {}}\n'
                 '{"jsonrpc": "2.0", "id": 2, "method": "tools/list"}\n'
                 '{"jsonrpc": "2.0", "id": 3, "method": "ping"}\n'
                 '{"jsonrpc": "2.0", "method": "notifications/initialized"}\n'
                 '\n'
                 '{"jsonrpc": "2.0", "id": 4, "method": "bogus"}\n')
        out = io.StringIO()
        with mock.patch.object(sys, "stdin", io.StringIO(lines)), \
                mock.patch.object(sys, "stdout", out):
            ws_main()
        rows = [json.loads(x) for x in out.getvalue().splitlines()]
        self.assertEqual(len(rows), 5)
        self.assertEqual(rows[0]["error"]["code"], -32700)
        self.assertEqual(rows[1]["result"]["serverInfo"]["name"],
                         "websearch")
        self.assertEqual(rows[3]["result"], {})
        self.assertEqual(rows[4]["error"]["code"], -32601)


class ImageGenMcpTest(unittest.TestCase):
    def test_tools_call_and_errors(self):
        gen = lambda p, output_filename=None, size="1024x1024", model=None: {  # noqa: E731
            "path": "/w/gen.png", "prompt": p}
        edit = lambda p, inp, output_filename=None, size="1024x1024", model=None: {  # noqa: E731
            "path": "/w/edit.png", "input": inp}
        resp = ig_handle({"jsonrpc": "2.0", "id": 1, "method": "tools/call",
                          "params": {"name": "image_generate",
                                     "arguments": {"prompt": "cat"}}},
                         generate_fn=gen)
        self.assertIn("gen.png", resp["result"]["content"][0]["text"])
        resp = ig_handle({"jsonrpc": "2.0", "id": 2, "method": "tools/call",
                          "params": {"name": "image_edit",
                                     "arguments": {"prompt": "dog",
                                                   "input_filename": "in.png"}}},
                         edit_fn=edit)
        self.assertIn("in.png", resp["result"]["content"][0]["text"])
        unknown = _call(ig_handle, "tools/call", name="nope")
        self.assertEqual(unknown["error"]["code"], -32602)
        boom = ig_handle({"jsonrpc": "2.0", "id": 3, "method": "tools/call",
                          "params": {"name": "image_generate",
                                     "arguments": {"prompt": "x"}}},
                         generate_fn=lambda *a, **k: (_ for _ in ()).throw(
                             RuntimeError("gen boom")))
        self.assertTrue(boom["result"]["isError"])
        self.assertIn("gen boom", boom["result"]["content"][0]["text"])
        self.assertEqual(ig_handle(42)["error"]["code"], -32700)
        self.assertIsNone(ig_handle({"method": "ping"}))
        self.assertEqual(_call(ig_handle, "ping")["result"], {})

    def test_main_stdio_loop(self):
        lines = ('{"jsonrpc": "2.0", "id": 1, "method": "initialize",'
                 ' "params": {}}\n'
                 '{"jsonrpc": "2.0", "id": 2, "method": "tools/list"}\n'
                 'not json\n')
        out = io.StringIO()
        with mock.patch.object(sys, "stdin", io.StringIO(lines)), \
                mock.patch.object(sys, "stdout", out):
            ig_main()
        rows = [json.loads(x) for x in out.getvalue().splitlines()]
        self.assertEqual(len(rows), 3)
        self.assertEqual(rows[0]["result"]["serverInfo"]["name"],
                         "imagegen")
        self.assertEqual(rows[2]["error"]["code"], -32700)


class FakeBrowserFull:
    def __init__(self):
        self.calls = []

    def navigate(self, url):
        self.calls.append(("navigate", url))
        return {"tool": "navigate", "url": url}

    def snapshot(self):
        self.calls.append(("snapshot",))
        return {"tool": "snapshot"}

    def click(self, selector):
        self.calls.append(("click", selector))
        return {"tool": "click"}

    def fill(self, selector, text):
        self.calls.append(("fill", selector, text))
        return {"tool": "fill"}

    def click_at(self, x, y):
        self.calls.append(("click_at", x, y))
        return {"tool": "click_at"}

    def mouse_move(self, x, y):
        self.calls.append(("mouse_move", x, y))
        return {"tool": "mouse_move"}

    def key_press(self, key):
        self.calls.append(("key_press", key))
        return {"tool": "key_press"}

    def type_text(self, text):
        self.calls.append(("type_text", text))
        return {"tool": "type_text"}

    def snapshot_refs(self):
        self.calls.append(("snapshot_refs",))
        return {"tool": "snapshot_refs"}

    def click_ref(self, ref):
        self.calls.append(("click_ref", ref))
        return {"tool": "click_ref"}

    def live_shot(self, name=None):
        self.calls.append(("live_shot", name))
        return {"tool": "live_shot"}

    def evaluate(self, expression):
        self.calls.append(("evaluate", expression))
        return {"tool": "evaluate"}

    def screenshot(self, name=None):
        self.calls.append(("screenshot", name))
        return {"tool": "screenshot"}

    def close(self):
        self.calls.append(("close",))
        return {"tool": "close"}


class PlaywrightMcpTest(unittest.TestCase):
    def test_all_tools_dispatch(self):
        cases = [
            ("browser_navigate", {"url": "https://x/"}, "navigate"),
            ("browser_snapshot", {}, "snapshot"),
            ("browser_click", {"selector": "#a"}, "click"),
            ("browser_fill", {"selector": "#a", "text": "hi"}, "fill"),
            ("browser_click_at", {"x": 1, "y": 2}, "click_at"),
            ("browser_mouse_move", {"x": 1, "y": 2}, "mouse_move"),
            ("browser_key_press", {"key": "Enter"}, "key_press"),
            ("browser_type_text", {"text": "hi"}, "type_text"),
            ("browser_snapshot_refs", {}, "snapshot_refs"),
            ("browser_click_ref", {"ref": "e0"}, "click_ref"),
            ("browser_live_shot", {}, "live_shot"),
            ("browser_evaluate", {"expression": "1+1"}, "evaluate"),
            ("browser_screenshot", {}, "screenshot"),
            ("browser_close", {}, "close"),
        ]
        for i, (name, args, tool) in enumerate(cases):
            fake = FakeBrowserFull()
            resp = pw_handle(
                {"jsonrpc": "2.0", "id": i + 1, "method": "tools/call",
                 "params": {"name": name, "arguments": args}},
                browser=fake)
            self.assertFalse(resp["result"]["isError"], name)
            self.assertIn(tool, resp["result"]["content"][0]["text"], name)

    def test_unknown_tool_and_tool_exception(self):
        fake = FakeBrowserFull()
        unknown = pw_handle(
            {"jsonrpc": "2.0", "id": 1, "method": "tools/call",
             "params": {"name": "browser_nope", "arguments": {}}},
            browser=fake)
        self.assertEqual(unknown["error"]["code"], -32602)

        class Raising(FakeBrowserFull):
            def navigate(self, url):
                raise RuntimeError("nav boom")

        resp = pw_handle(
            {"jsonrpc": "2.0", "id": 2, "method": "tools/call",
             "params": {"name": "browser_navigate",
                        "arguments": {"url": "https://x/"}}},
            browser=Raising())
        self.assertTrue(resp["result"]["isError"])
        self.assertIn("nav boom", resp["result"]["content"][0]["text"])

    def test_protocol_and_handle_line(self):
        self.assertEqual(pw_handle("x")["error"]["code"], -32700)
        self.assertIsNone(pw_handle({"method": "ping"}))
        self.assertEqual(_call(pw_handle, "ping")["result"], {})
        out = io.StringIO()
        with mock.patch.object(sys, "stdout", out):
            _handle_line('{"jsonrpc": "2.0", "id": 9, "method": "ping"}',
                         FakeBrowserFull())
            _handle_line("{bad json", FakeBrowserFull())
        rows = [json.loads(x) for x in out.getvalue().splitlines()]
        self.assertEqual(rows[0]["result"], {})
        self.assertEqual(rows[1]["error"]["code"], -32700)
        with mock.patch.object(pw_mod, "handle_request",
                               side_effect=RuntimeError("handler boom")):
            out2, err2 = io.StringIO(), io.StringIO()
            with mock.patch.object(sys, "stdout", out2), \
                    mock.patch.object(sys, "stderr", err2):
                _handle_line('{"jsonrpc": "2.0", "id": 10, '
                             '"method": "tools/list"}',
                             FakeBrowserFull())
        row = json.loads(out2.getvalue().strip())
        self.assertEqual(row["error"]["code"], -32603)
        self.assertIn("handler boom", err2.getvalue())


class FlatCopyTest(unittest.TestCase):
    def _flat_copy(self, pairs):
        td = tempfile.mkdtemp()
        self.addCleanup(shutil.rmtree, td, True)
        root = Path(__file__).resolve().parents[2] / "cloud"
        for name in pairs:
            shutil.copy(root / name, os.path.join(td, name))
        return td

    def _run_flat(self, td, module, check):
        script = ("import %s as m; print(m.SERVER_NAME); print(check(m))"
                  % module)
        env = {"PATH": os.environ.get("PATH", ""), "PYTHONPATH": td}
        proc = subprocess.run(
            [sys.executable, "-c",
             script.replace("check(m)", check)],
            cwd=td, env=env, capture_output=True, text=True, timeout=60)
        self.assertEqual(proc.returncode, 0, proc.stderr)
        return proc.stdout

    def test_playwright_flat_copy(self):
        td = self._flat_copy(["playwright_mcp.py",
                              "playwright_browser.py"])
        out = self._run_flat(
            td, "playwright_mcp",
            "m.Browser.__module__ + '|' + "
            "str(m.handle_request({'jsonrpc': '2.0', 'id': 1,"
            " 'method': 'tools/list'})['result']['tools'][0]['name'])")
        self.assertIn("playwright", out)
        self.assertIn("playwright_browser|browser_navigate", out)

    def test_websearch_flat_copy(self):
        td = self._flat_copy(["websearch_mcp.py", "websearch.py"])
        out = self._run_flat(
            td, "websearch_mcp",
            "m.search.__module__ + '|' + "
            "str(m.handle_request({'jsonrpc': '2.0', 'id': 1,"
            " 'method': 'ping'})['result'])")
        self.assertIn("websearch", out)
        self.assertIn("websearch|{}", out)

    def test_imagegen_flat_copy(self):
        td = self._flat_copy(["imagegen_mcp.py", "imagegen.py"])
        out = self._run_flat(
            td, "imagegen_mcp",
            "m.generate.__module__ + '|' + "
            "str(m.handle_request({'jsonrpc': '2.0', 'id': 1,"
            " 'method': 'ping'})['result'])")
        self.assertIn("imagegen", out)
        self.assertIn("imagegen|{}", out)


if __name__ == "__main__":
    unittest.main()