"""Map the editor ``browser_*`` tool names to ``BrowserTools`` methods and
serialize results exactly as the pod's MCP does.

This is the small glue that mirrors ``MainController.browser_*``: it forwards a
tool call to the right ``BrowserTools`` method (positionally — the method param
names differ from the tool arg names, e.g. ``page_id`` vs ``tab_id``, so we pass
by position like the pod does) and runs the result through the real
``serialize_tool_result`` so cloud-api receives byte-identical output.

Screenshot is special-cased: ``BrowserTools.screenshot`` returns a temp file
Path on THIS pod, so we read+base64 it into the ``__imageContent`` dict here
(exactly as ``MainController.browser_screenshot`` does) before serializing.
"""

import base64
from pathlib import Path
from typing import Any, Dict, List, Optional, Protocol, Tuple

from abstra_internals.utils.mcp import serialize_tool_result

_REQUIRED = object()

# Mirrors main.py EDITOR_PAGE_NAV_TIMEOUT_MS.
_PAGE_NAV_TIMEOUT_MS = 150_000

# tool name -> (BrowserTools method, [(arg_name, default), ...]) in positional
# order. Derived from MainController.browser_* method bodies.
_TOOLS: Dict[str, Tuple[str, List[Tuple[str, Any]]]] = {
    "browser_navigate": ("navigate_to_url", [("url", _REQUIRED)]),
    "browser_open_page": (
        "navigate_to_url",
        [
            ("url", _REQUIRED),
            ("page_id", None),
            ("wait_until", "domcontentloaded"),
            ("timeout_ms", _PAGE_NAV_TIMEOUT_MS),
        ],
    ),
    "browser_get_page_summary": ("get_page_summary", [("tab_id", _REQUIRED)]),
    "browser_click": ("click_element", [("tab_id", _REQUIRED), ("index", _REQUIRED)]),
    "browser_fill": (
        "fill_element",
        [("tab_id", _REQUIRED), ("index", _REQUIRED), ("value", _REQUIRED)],
    ),
    "browser_get_text": ("get_text", [("tab_id", _REQUIRED), ("selector", _REQUIRED)]),
    "browser_get_html": (
        "get_element_html",
        [("tab_id", _REQUIRED), ("selector", _REQUIRED)],
    ),
    "browser_execute_javascript": (
        "execute_javascript",
        [("tab_id", _REQUIRED), ("script", _REQUIRED)],
    ),
    "browser_wait": ("wait", [("tab_id", None), ("milliseconds", 1000)]),
    "browser_get_console_logs": ("get_console_logs", [("tab_id", _REQUIRED)]),
    "browser_get_network_requests": ("get_network_requests", [("tab_id", _REQUIRED)]),
    "browser_close": ("close_page", [("tab_id", _REQUIRED)]),
    "browser_list_tabs": ("list_pages", []),
}

SCREENSHOT_TOOL = "browser_screenshot"

# The 14 tools the pool serves (matches cloud-api's
# inProcessEditorToolNames.browserAutomation).
BROWSER_SERVICE_TOOL_NAMES: Tuple[str, ...] = tuple(_TOOLS.keys()) + (SCREENSHOT_TOOL,)


class SessionLike(Protocol):
    def call(self, method_name: str, *args: Any) -> Any: ...


class UnknownBrowserTool(Exception):
    pass


class MissingToolArgument(Exception):
    pass


def _positional_args(
    tool_name: str, specs: List[Tuple[str, Any]], args: Dict[str, Any]
) -> List[Any]:
    out: List[Any] = []
    for name, default in specs:
        if name in args and args[name] is not None:
            out.append(args[name])
        elif default is _REQUIRED:
            raise MissingToolArgument(f"{tool_name} requires argument '{name}'")
        else:
            out.append(default)
    return out


def dispatch(
    session: SessionLike, tool_name: str, args: Optional[Dict[str, Any]]
) -> str:
    """Run one editor browser tool against ``session`` and return the
    pod-identical serialized string."""
    args = args or {}
    if tool_name == SCREENSHOT_TOOL:
        return _screenshot(session, args)
    if tool_name not in _TOOLS:
        raise UnknownBrowserTool(tool_name)
    method, specs = _TOOLS[tool_name]
    positional = _positional_args(tool_name, specs, args)
    return serialize_tool_result(session.call(method, *positional))


def _screenshot(session: SessionLike, args: Dict[str, Any]) -> str:
    tab_id = args.get("tab_id")
    if tab_id is None:
        raise MissingToolArgument("browser_screenshot requires argument 'tab_id'")
    show_markers = bool(args.get("show_markers", False))
    highlight_element = args.get("highlight_element")

    image_path = Path(
        session.call("screenshot", tab_id, show_markers, highlight_element)
    )
    try:
        image_bytes = image_path.read_bytes()
    finally:
        image_path.unlink(missing_ok=True)

    encoded = base64.b64encode(image_bytes).decode("ascii")
    details = []
    if show_markers:
        details.append("with interactive-element markers")
    if highlight_element:
        details.append(f"highlighting {highlight_element!r}")
    suffix = f" ({', '.join(details)})" if details else ""
    # __imageContent is the cloud-api toolCaller contract; serialize_tool_result
    # JSON-stringifies it (a dict), and cloud-api parses it back to attach the
    # image — identical to the pod path.
    payload = {
        "__imageContent": True,
        "mimeType": "image/jpeg",
        "dataUri": f"data:image/jpeg;base64,{encoded}",
        "sizeKb": len(image_bytes) // 1024,
        "description": f"Screenshot of tab {tab_id}{suffix}",
    }
    return serialize_tool_result(payload)
