"""Interactive browser tools backed by Playwright or Playwright MCP."""

from __future__ import annotations

import asyncio
import json
import os
import shutil
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from loguru import logger

from madpod_agent.agent.tools.base import Tool
from madpod_agent.agent.tools.registry import ToolRegistry
from madpod_agent.config.schema import BrowserToolsConfig

_REF_ATTR = "data-nanobot-ref"


def _json(payload: dict[str, Any]) -> str:
    return json.dumps(payload, ensure_ascii=False)


def _maybe_int(value: Any, default: int) -> int:
    try:
        return int(value)
    except (TypeError, ValueError):
        return default


@dataclass
class BrowserSnapshotState:
    """State cached from the latest browser snapshot."""

    refs: dict[str, str] = field(default_factory=dict)


class BrowserBackend:
    """Backend contract for interactive browser tools."""

    async def navigate(self, *, url: str) -> str:
        raise NotImplementedError

    async def snapshot(self, *, max_elements: int | None = None) -> str:
        raise NotImplementedError

    async def click(self, **kwargs: Any) -> str:
        raise NotImplementedError

    async def type(self, **kwargs: Any) -> str:
        raise NotImplementedError

    async def wait(self, **kwargs: Any) -> str:
        raise NotImplementedError

    async def back(self) -> str:
        raise NotImplementedError

    async def take_screenshot(self, **kwargs: Any) -> str:
        raise NotImplementedError

    async def execute(self, **kwargs: Any) -> str:
        raise NotImplementedError


class NativePlaywrightBrowserBackend(BrowserBackend):
    """Native Playwright-backed browser automation."""

    def __init__(self, workspace: Path, config: BrowserToolsConfig):
        self.workspace = workspace
        self.config = config
        self._lock = asyncio.Lock()
        self._playwright = None
        self._context = None
        self._page = None
        self._state = BrowserSnapshotState()

    def _profile_dir(self) -> Path:
        return self.workspace / self.config.profile_subdir

    def _output_dir(self) -> Path:
        return self.workspace / self.config.output_subdir

    def _headless(self) -> bool:
        if self.config.headless is not None:
            return self.config.headless
        return not bool(os.environ.get("DISPLAY"))

    @staticmethod
    def _normalize_url(url: str) -> str:
        value = (url or "").strip()
        if not value:
            return value
        if "://" not in value:
            return "https://" + value
        return value

    @staticmethod
    def _find_browser_binary() -> str | None:
        for candidate in (
            "google-chrome",
            "google-chrome-stable",
            "chrome",
            "chromium-browser",
            "chromium",
        ):
            path = shutil.which(candidate)
            if path:
                return path
        return None

    async def _ensure_page(self):
        async with self._lock:
            if self._page and not self._page.is_closed():
                return self._page

            try:
                from playwright.async_api import async_playwright
            except Exception as exc:
                raise RuntimeError(
                    "Playwright is not installed in this runtime. "
                    "Install the 'playwright' dependency or switch tools.browser.backend to 'playwright_mcp'."
                ) from exc

            if self._playwright is None:
                self._playwright = await async_playwright().start()

            self._profile_dir().mkdir(parents=True, exist_ok=True)
            self._output_dir().mkdir(parents=True, exist_ok=True)

            launch_kwargs: dict[str, Any] = {
                "user_data_dir": str(self._profile_dir()),
                "headless": self._headless(),
                "viewport": {
                    "width": self.config.viewport_width,
                    "height": self.config.viewport_height,
                },
                "args": ["--no-sandbox"],
            }
            if executable_path := self._find_browser_binary():
                launch_kwargs["executable_path"] = executable_path

            self._context = await self._playwright.chromium.launch_persistent_context(**launch_kwargs)
            self._context.set_default_timeout(self.config.action_timeout_s * 1000)
            self._context.set_default_navigation_timeout(self.config.navigation_timeout_s * 1000)
            self._page = self._context.pages[0] if self._context.pages else await self._context.new_page()
            try:
                await self._page.bring_to_front()
            except Exception:
                pass
            return self._page

    async def _current_page(self):
        page = await self._ensure_page()
        if self._context:
            for candidate in self._context.pages:
                if not candidate.is_closed():
                    self._page = candidate
                    break
        return self._page or page

    async def _tabs(self) -> list[dict[str, Any]]:
        if not self._context:
            return []
        active = await self._current_page()
        tabs = []
        for index, page in enumerate(self._context.pages):
            if page.is_closed():
                continue
            tabs.append(
                {
                    "index": index,
                    "active": page == active,
                    "url": page.url,
                    "title": await page.title(),
                }
            )
        return tabs

    async def navigate(self, *, url: str) -> str:
        target = self._normalize_url(url)
        if not target:
            return "Error: 'url' is required"

        page = await self._current_page()
        await page.goto(target, wait_until="domcontentloaded")
        return _json(
            {
                "ok": True,
                "action": "navigate",
                "url": page.url,
                "title": await page.title(),
                "headless": self._headless(),
                "message": "Browser opened. Run browser_snapshot to inspect the page.",
            }
        )

    async def snapshot(self, *, max_elements: int | None = None) -> str:
        page = await self._current_page()
        limit = max(1, min(max_elements or self.config.snapshot_max_elements, 100))
        payload = await page.evaluate(
            """
            ({ attrName, maxElements, textChars }) => {
              const root = document.querySelector('main') || document.body || document.documentElement;
              const compact = (value, max = 160) => {
                const text = String(value || '').replace(/\\s+/g, ' ').trim();
                return text.length > max ? text.slice(0, max - 1) + '…' : text;
              };
              const visible = (el) => {
                if (!(el instanceof HTMLElement)) return false;
                const style = window.getComputedStyle(el);
                if (!style || style.visibility === 'hidden' || style.display === 'none') return false;
                const rect = el.getBoundingClientRect();
                if (rect.width < 2 || rect.height < 2) return false;
                return rect.bottom >= 0 && rect.right >= 0
                  && rect.top <= window.innerHeight && rect.left <= window.innerWidth;
              };
              const labelFor = (el) => compact(
                el.getAttribute('aria-label')
                || el.getAttribute('placeholder')
                || el.getAttribute('title')
                || el.innerText
                || el.textContent
                || el.getAttribute('name')
                || el.getAttribute('value')
                || el.getAttribute('href')
              );
              window.__nanobotRefSeq = window.__nanobotRefSeq || 0;
              const selectors = [
                'a[href]',
                'button',
                'input',
                'textarea',
                'select',
                '[role="button"]',
                '[role="link"]',
                '[contenteditable="true"]',
                '[tabindex]'
              ];
              const seen = new Set();
              const elements = [];
              for (const el of document.querySelectorAll(selectors.join(','))) {
                if (elements.length >= maxElements) break;
                if (!visible(el) || seen.has(el)) continue;
                seen.add(el);
                let ref = el.getAttribute(attrName);
                if (!ref) {
                  ref = `e${++window.__nanobotRefSeq}`;
                  el.setAttribute(attrName, ref);
                }
                const rect = el.getBoundingClientRect();
                elements.push({
                  ref: `@${ref}`,
                  tag: (el.tagName || '').toLowerCase(),
                  role: el.getAttribute('role') || '',
                  type: el.getAttribute('type') || '',
                  label: labelFor(el),
                  text: compact(el.innerText || el.textContent || '', 120),
                  href: el.getAttribute('href') || '',
                  x: Math.round(rect.x),
                  y: Math.round(rect.y),
                });
              }
              const textSource = compact(root ? root.innerText || root.textContent || '' : '', textChars);
              return {
                url: location.href,
                title: document.title || '',
                elements,
                textExcerpt: textSource,
                note: 'Refs can change after navigation or DOM updates. Re-run browser_snapshot after the page changes.',
              };
            }
            """,
            {
                "attrName": _REF_ATTR,
                "maxElements": limit,
                "textChars": self.config.snapshot_text_chars,
            },
        )
        self._state.refs = {
            item["ref"]: f'[{_REF_ATTR}="{item["ref"].lstrip("@")}"]'
            for item in payload.get("elements", [])
        }
        payload["tabs"] = await self._tabs()
        return _json(payload)

    def _target_selector(self, *, ref: str | None = None, selector: str | None = None) -> str:
        if ref:
            key = ref if ref.startswith("@") else f"@{ref}"
            return self._state.refs.get(key) or f'[{_REF_ATTR}="{key.lstrip("@")}"]'
        if selector:
            return selector
        raise ValueError("Provide 'ref' or 'selector'")

    async def click(self, **kwargs: Any) -> str:
        page = await self._current_page()
        target = self._target_selector(ref=kwargs.get("ref"), selector=kwargs.get("selector"))
        locator = page.locator(target).first
        await locator.click(button=str(kwargs.get("button") or "left"))
        try:
            await page.wait_for_load_state("domcontentloaded", timeout=self.config.navigation_timeout_s * 1000)
        except Exception:
            pass
        return _json(
            {
                "ok": True,
                "action": "click",
                "url": page.url,
                "title": await page.title(),
                "target": kwargs.get("ref") or kwargs.get("selector"),
                "message": "Click completed. Re-run browser_snapshot if the page changed.",
            }
        )

    async def type(self, **kwargs: Any) -> str:
        text = str(kwargs.get("text") or "")
        if not text:
            return "Error: 'text' is required"
        page = await self._current_page()
        target = self._target_selector(ref=kwargs.get("ref"), selector=kwargs.get("selector"))
        locator = page.locator(target).first
        if kwargs.get("clear_first", True):
            await locator.fill(text)
        else:
            await locator.type(text)
        if kwargs.get("submit"):
            await locator.press("Enter")
        return _json(
            {
                "ok": True,
                "action": "type",
                "url": page.url,
                "title": await page.title(),
                "target": kwargs.get("ref") or kwargs.get("selector"),
            }
        )

    async def wait(self, **kwargs: Any) -> str:
        page = await self._current_page()
        if kwargs.get("seconds") is not None:
            await asyncio.sleep(max(0, float(kwargs["seconds"])))
        elif kwargs.get("text"):
            await page.get_by_text(str(kwargs["text"])).first.wait_for(
                timeout=self.config.navigation_timeout_s * 1000
            )
        elif kwargs.get("text_gone"):
            await page.wait_for_function(
                """value => !document.body || !document.body.innerText.includes(value)""",
                str(kwargs["text_gone"]),
                timeout=self.config.navigation_timeout_s * 1000,
            )
        elif kwargs.get("ref") or kwargs.get("selector"):
            target = self._target_selector(ref=kwargs.get("ref"), selector=kwargs.get("selector"))
            await page.locator(target).first.wait_for(timeout=self.config.navigation_timeout_s * 1000)
        else:
            await asyncio.sleep(1.0)
        return _json({"ok": True, "action": "wait", "url": page.url, "title": await page.title()})

    async def back(self) -> str:
        page = await self._current_page()
        await page.go_back(wait_until="domcontentloaded")
        return _json({"ok": True, "action": "back", "url": page.url, "title": await page.title()})

    async def take_screenshot(self, **kwargs: Any) -> str:
        page = await self._current_page()
        output_dir = self._output_dir()
        output_dir.mkdir(parents=True, exist_ok=True)
        filename = str(kwargs.get("filename") or f"screenshot-{int(time.time())}.png")
        target = output_dir / Path(filename).name
        await page.screenshot(path=str(target), full_page=bool(kwargs.get("full_page")))
        return _json(
            {
                "ok": True,
                "action": "screenshot",
                "url": page.url,
                "title": await page.title(),
                "path": str(target),
            }
        )

    async def execute(self, **kwargs: Any) -> str:
        code = str(kwargs.get("code") or "").strip()
        if not code:
            return "Error: 'code' is required"
        page = await self._current_page()
        result = await page.evaluate(code)
        return _json(
            {
                "ok": True,
                "action": "execute",
                "url": page.url,
                "title": await page.title(),
                "result": result if isinstance(result, (str, int, float, bool, list, dict)) else str(result),
            }
        )


class MCPBrowserBackend(BrowserBackend):
    """Stable browser tool surface backed by hidden Playwright MCP tools."""

    TOOL_MAP = {
        "browser_navigate": "browser_navigate",
        "browser_snapshot": "browser_snapshot",
        "browser_click": "browser_click",
        "browser_type": "browser_type",
        "browser_wait": "browser_wait_for",
        "browser_back": "browser_navigate_back",
        "browser_take_screenshot": "browser_take_screenshot",
        "browser_execute": "browser_evaluate",
    }

    def __init__(self, registry: ToolRegistry, server_name: str):
        self.registry = registry
        self.server_name = server_name

    async def _call(self, stable_name: str, params: dict[str, Any]) -> str:
        mapped = self.TOOL_MAP[stable_name]
        tool_name = f"mcp_{self.server_name}_{mapped}"
        if not self.registry.has(tool_name):
            return (
                f"Error: MCP browser backend is enabled but '{tool_name}' is unavailable. "
                "Ensure the Playwright MCP server is configured and connected."
            )
        return await self.registry.execute(tool_name, params)

    async def navigate(self, *, url: str) -> str:
        return await self._call("browser_navigate", {"url": url})

    async def snapshot(self, *, max_elements: int | None = None) -> str:
        params = {}
        if max_elements is not None:
            params["maxElements"] = max_elements
        return await self._call("browser_snapshot", params)

    async def click(self, **kwargs: Any) -> str:
        params = {
            "element": kwargs.get("ref") or kwargs.get("selector") or "",
            "ref": kwargs.get("ref") or "",
        }
        return await self._call("browser_click", params)

    async def type(self, **kwargs: Any) -> str:
        params = {
            "element": kwargs.get("ref") or kwargs.get("selector") or "",
            "ref": kwargs.get("ref") or "",
            "text": kwargs.get("text") or "",
            "submit": bool(kwargs.get("submit")),
        }
        return await self._call("browser_type", params)

    async def wait(self, **kwargs: Any) -> str:
        params: dict[str, Any] = {}
        if kwargs.get("seconds") is not None:
            params["time"] = kwargs.get("seconds")
        if kwargs.get("text"):
            params["text"] = kwargs.get("text")
        if kwargs.get("text_gone"):
            params["textGone"] = kwargs.get("text_gone")
        return await self._call("browser_wait", params)

    async def back(self) -> str:
        return await self._call("browser_back", {})

    async def take_screenshot(self, **kwargs: Any) -> str:
        params = {}
        if kwargs.get("filename"):
            params["filename"] = kwargs.get("filename")
        if kwargs.get("full_page") is not None:
            params["fullPage"] = kwargs.get("full_page")
        return await self._call("browser_take_screenshot", params)

    async def execute(self, **kwargs: Any) -> str:
        return await self._call("browser_execute", {"function": kwargs.get("code") or ""})


class BrowserToolBase(Tool):
    """Shared base for first-class browser tools."""

    def __init__(self, backend: BrowserBackend):
        self.backend = backend


class BrowserNavigateTool(BrowserToolBase):
    @property
    def name(self) -> str:
        return "browser_navigate"

    @property
    def description(self) -> str:
        return (
            "Open Chrome/Chromium and navigate to a URL. "
            "Use this for interactive, dynamic, or authenticated sites where web_search/web_fetch cannot see the real page state."
        )

    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "url": {"type": "string", "description": "The URL to open"},
            },
            "required": ["url"],
        }

    async def execute(self, **kwargs: Any) -> str:
        return await self.backend.navigate(url=kwargs["url"])


class BrowserSnapshotTool(BrowserToolBase):
    @property
    def name(self) -> str:
        return "browser_snapshot"

    @property
    def description(self) -> str:
        return (
            "Return a compact snapshot of the current page with reusable element refs. "
            "Use this after navigation and after visible page changes so later click/type actions target the current DOM instead of stale refs."
        )

    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "max_elements": {
                    "type": "integer",
                    "minimum": 1,
                    "maximum": 100,
                    "description": "Optional limit for returned interactive elements",
                }
            },
        }

    async def execute(self, **kwargs: Any) -> str:
        return await self.backend.snapshot(max_elements=kwargs.get("max_elements"))


class BrowserClickTool(BrowserToolBase):
    @property
    def name(self) -> str:
        return "browser_click"

    @property
    def description(self) -> str:
        return (
            "Click an element in the current browser page using a snapshot ref or selector. "
            "Prefer refs from browser_snapshot because they are more stable and easier for the model to reason about than raw selectors."
        )

    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "ref": {"type": "string", "description": "Element ref from browser_snapshot, e.g. @e1"},
                "selector": {"type": "string", "description": "CSS selector fallback"},
                "button": {"type": "string", "enum": ["left", "right", "middle"]},
            },
        }

    async def execute(self, **kwargs: Any) -> str:
        return await self.backend.click(**kwargs)


class BrowserTypeTool(BrowserToolBase):
    @property
    def name(self) -> str:
        return "browser_type"

    @property
    def description(self) -> str:
        return (
            "Type into a field in the current browser page using a snapshot ref or selector. "
            "Use refs from browser_snapshot when possible, and re-snapshot first if navigation or DOM changes may have invalidated prior refs."
        )

    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "ref": {"type": "string", "description": "Element ref from browser_snapshot, e.g. @e1"},
                "selector": {"type": "string", "description": "CSS selector fallback"},
                "text": {"type": "string", "description": "Text to type"},
                "clear_first": {"type": "boolean", "description": "Clear the field before typing"},
                "submit": {"type": "boolean", "description": "Press Enter after typing"},
            },
            "required": ["text"],
        }

    async def execute(self, **kwargs: Any) -> str:
        return await self.backend.type(**kwargs)


class BrowserWaitTool(BrowserToolBase):
    @property
    def name(self) -> str:
        return "browser_wait"

    @property
    def description(self) -> str:
        return (
            "Wait for time, text, or an element while the browser page updates. "
            "Use this when a click, form submit, or navigation triggers async UI changes before the next snapshot or action."
        )

    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "seconds": {"type": "number", "minimum": 0},
                "ref": {"type": "string"},
                "selector": {"type": "string"},
                "text": {"type": "string"},
                "text_gone": {"type": "string"},
            },
        }

    async def execute(self, **kwargs: Any) -> str:
        return await self.backend.wait(**kwargs)


class BrowserBackTool(BrowserToolBase):
    @property
    def name(self) -> str:
        return "browser_back"

    @property
    def description(self) -> str:
        return "Go back one step in browser history."

    @property
    def parameters(self) -> dict[str, Any]:
        return {"type": "object", "properties": {}}

    async def execute(self, **kwargs: Any) -> str:
        return await self.backend.back()


class BrowserTakeScreenshotTool(BrowserToolBase):
    @property
    def name(self) -> str:
        return "browser_take_screenshot"

    @property
    def description(self) -> str:
        return (
            "Capture a screenshot of the current browser page. "
            "Use this when visual verification matters or when you need to inspect a layout/state that text extraction may miss."
        )

    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "filename": {"type": "string"},
                "full_page": {"type": "boolean"},
            },
        }

    async def execute(self, **kwargs: Any) -> str:
        return await self.backend.take_screenshot(**kwargs)


class BrowserExecuteTool(BrowserToolBase):
    @property
    def name(self) -> str:
        return "browser_execute"

    @property
    def description(self) -> str:
        return (
            "Run JavaScript in the current browser page and return the result. "
            "Use this for DOM inspection or state extraction when snapshot text is insufficient, not as a first-choice replacement for normal browser tools."
        )

    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "code": {"type": "string", "description": "JavaScript expression or function body"},
            },
            "required": ["code"],
        }

    async def execute(self, **kwargs: Any) -> str:
        return await self.backend.execute(**kwargs)


def register_browser_tools(
    registry: ToolRegistry,
    *,
    workspace: Path,
    config: BrowserToolsConfig | None,
) -> None:
    """Register first-class browser tools when enabled."""
    if config is None or not config.enabled:
        return

    backend: BrowserBackend
    if config.backend == "playwright_mcp":
        backend = MCPBrowserBackend(registry, config.mcp_server_name)
    else:
        backend = NativePlaywrightBrowserBackend(workspace, config)

    for tool_cls in (
        BrowserNavigateTool,
        BrowserSnapshotTool,
        BrowserClickTool,
        BrowserTypeTool,
        BrowserWaitTool,
        BrowserBackTool,
        BrowserTakeScreenshotTool,
        BrowserExecuteTool,
    ):
        tool = tool_cls(backend)
        registry.register(tool)
        logger.debug("Registered browser tool '{}'", tool.name)
