#!/usr/bin/env python3
"""Browser MCP v2 — Playwright Chromium + content cache + resource blocking."""
import asyncio, base64, json, os, subprocess, shutil, threading
from pathlib import Path
from typing import Optional

from cachetools import TTLCache

SCRIPTS_DIR = Path("/scripts")
SCRIPTS_DIR.mkdir(exist_ok=True)

# Content cache: URL → text/html (5-min TTL, max 128 pages)
_page_cache: TTLCache = TTLCache(maxsize=128, ttl=300)
_cache_lock = asyncio.Lock()

import mcp.types as types
from mcp.server import Server
from mcp.server.sse import SseServerTransport
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import Response
from starlette.routing import Mount, Route
import uvicorn
from playwright.async_api import async_playwright, Browser, BrowserContext, Page

mcp_server = Server("browser-mcp")
sse = SseServerTransport("/browser/messages/")

# Optimized Chrome flags — sourced from Playwright Docker + headless benchmarks
_CHROME_FLAGS = [
    "--headless=new",
    "--no-sandbox",
    "--disable-setuid-sandbox",
    "--disable-dev-shm-usage",          # critical in Docker: /dev/shm is only 64 MB
    "--disable-gpu",
    "--disable-software-rasterizer",
    "--disable-extensions",
    "--disable-background-networking",   # stops background DNS/resource prefetch
    "--disable-backgrounding-occluded-windows",
    "--disable-renderer-backgrounding",
    "--disable-background-timer-throttling",
    "--no-first-run",
    "--no-zygote",
    "--mute-audio",
    "--hide-scrollbars",
    "--window-size=1280,720",
]

# Resource types to block when fast_mode=True on navigate
_BLOCK_TYPES = {"image", "media", "font", "stylesheet"}


class BrowserManager:
    def __init__(self):
        self.pw = None
        self.browser: Optional[Browser] = None
        self.context: Optional[BrowserContext] = None
        self.tabs: list[Page] = []
        self.current: int = 0
        self.console_msgs: list[str] = []
        self._fast_mode: bool = False  # block images/fonts/css
        self._ready = threading.Event()

    async def init(self):
        self.pw = await async_playwright().start()
        self.browser = await self.pw.chromium.launch(
            headless=True,
            args=_CHROME_FLAGS,
        )
        self.context = await self.browser.new_context(
            viewport={"width": 1280, "height": 800},
            user_agent=(
                "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
                "(KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36"
            ),
        )
        await self._setup_routing(self.context)
        page = await self.context.new_page()
        page.on("console", lambda m: self.console_msgs.append(f"[{m.type}] {m.text}"))
        self.tabs = [page]
        self.current = 0
        self._ready.set()

    async def _setup_routing(self, ctx: BrowserContext):
        async def route_handler(route):
            if self._fast_mode and route.request.resource_type in _BLOCK_TYPES:
                await route.abort()
            else:
                await route.continue_()
        await ctx.route("**/*", route_handler)

    async def page(self) -> Page:
        if not self.tabs:
            await self.init()
        return self.tabs[self.current]

    async def new_tab(self, url: str = "about:blank") -> Page:
        p = await self.context.new_page()
        p.on("console", lambda m: self.console_msgs.append(f"[{m.type}] {m.text}"))
        if url and url != "about:blank":
            await p.goto(url, timeout=30000, wait_until="domcontentloaded")
        self.tabs.append(p)
        self.current = len(self.tabs) - 1
        return p


bm = BrowserManager()


def ok(t):
    return [types.TextContent(type="text", text=str(t))]


def img(data: bytes):
    return [types.ImageContent(type="image", data=base64.b64encode(data).decode(), mimeType="image/png")]


@mcp_server.list_tools()
async def list_tools():
    T = types.Tool
    return [
        T(name="navigate",
          description=(
              "Navigate to a URL. Default wait: domcontentloaded (fast). "
              "Use networkidle only if page needs JS to fully render. "
              "fast_mode=true blocks images/fonts/CSS — 30-50% faster for text extraction."
          ),
          inputSchema={"type": "object", "properties": {
              "url": {"type": "string"},
              "wait_until": {"type": "string", "default": "domcontentloaded",
                             "enum": ["load", "domcontentloaded", "networkidle"]},
              "fast_mode": {"type": "boolean", "default": False,
                            "description": "Block images/fonts/CSS to speed up text-only tasks"},
          }, "required": ["url"]}),

        T(name="screenshot",
          description="Take a screenshot of the current page or a specific element. Returns a PNG image.",
          inputSchema={"type": "object", "properties": {
              "full_page": {"type": "boolean", "default": False},
              "selector": {"type": "string", "description": "CSS selector to screenshot just that element"}
          }, "required": []}),

        T(name="get_text",
          description="Get visible text of the page or element. Cached 5 min per URL.",
          inputSchema={"type": "object", "properties": {
              "selector": {"type": "string", "default": "body"},
              "no_cache": {"type": "boolean", "default": False},
          }, "required": []}),

        T(name="get_html",
          description="Get HTML source of the page or element. Cached 5 min per URL.",
          inputSchema={"type": "object", "properties": {
              "selector": {"type": "string", "default": "html"},
              "outer": {"type": "boolean", "default": True},
              "no_cache": {"type": "boolean", "default": False},
          }, "required": []}),

        T(name="click",
          description="Click an element by CSS selector.",
          inputSchema={"type": "object", "properties": {
              "selector": {"type": "string"},
              "button": {"type": "string", "default": "left", "enum": ["left", "right", "middle"]},
              "count": {"type": "integer", "default": 1}
          }, "required": ["selector"]}),

        T(name="type_text",
          description="Type text into an input element.",
          inputSchema={"type": "object", "properties": {
              "selector": {"type": "string"},
              "text": {"type": "string"},
              "clear": {"type": "boolean", "default": False}
          }, "required": ["selector", "text"]}),

        T(name="fill",
          description="Fill a form field (clears first then sets value).",
          inputSchema={"type": "object", "properties": {
              "selector": {"type": "string"},
              "value": {"type": "string"}
          }, "required": ["selector", "value"]}),

        T(name="select_option",
          description="Select an option in a <select> dropdown.",
          inputSchema={"type": "object", "properties": {
              "selector": {"type": "string"},
              "value": {"type": "string"},
              "label": {"type": "string"}
          }, "required": ["selector"]}),

        T(name="hover",
          description="Hover over an element (triggers hover CSS/JS).",
          inputSchema={"type": "object", "properties": {
              "selector": {"type": "string"}
          }, "required": ["selector"]}),

        T(name="press_key",
          description="Press a keyboard key (Enter, Tab, Escape, ArrowDown, F5, etc.).",
          inputSchema={"type": "object", "properties": {
              "key": {"type": "string"},
              "selector": {"type": "string", "description": "Focus this element first (optional)"}
          }, "required": ["key"]}),

        T(name="evaluate",
          description="Execute JavaScript in the page and return the result.",
          inputSchema={"type": "object", "properties": {
              "script": {"type": "string"},
              "arg": {"type": "string", "description": "JSON argument passed as first param"}
          }, "required": ["script"]}),

        T(name="wait_for",
          description="Wait for an element to appear, or wait N milliseconds.",
          inputSchema={"type": "object", "properties": {
              "selector": {"type": "string"},
              "ms": {"type": "integer"},
              "state": {"type": "string", "default": "visible",
                       "enum": ["attached", "detached", "visible", "hidden"]}
          }, "required": []}),

        T(name="scroll",
          description="Scroll the page by pixels, or scroll an element into view.",
          inputSchema={"type": "object", "properties": {
              "x": {"type": "integer", "default": 0},
              "y": {"type": "integer", "default": 500},
              "selector": {"type": "string"}
          }, "required": []}),

        T(name="go_back",  description="Navigate back in history.",
          inputSchema={"type": "object", "properties": {}, "required": []}),
        T(name="go_forward", description="Navigate forward in history.",
          inputSchema={"type": "object", "properties": {}, "required": []}),
        T(name="reload",   description="Reload the current page.",
          inputSchema={"type": "object", "properties": {}, "required": []}),
        T(name="get_url",  description="Get current page URL.",
          inputSchema={"type": "object", "properties": {}, "required": []}),
        T(name="get_title",description="Get current page title.",
          inputSchema={"type": "object", "properties": {}, "required": []}),

        T(name="get_links",
          description="Get all links on the page (href + text). Cached 5 min.",
          inputSchema={"type": "object", "properties": {
              "selector": {"type": "string", "default": "a"},
              "limit": {"type": "integer", "default": 50},
              "no_cache": {"type": "boolean", "default": False},
          }, "required": []}),

        T(name="find_elements",
          description="Find elements by CSS selector, return their text or an attribute.",
          inputSchema={"type": "object", "properties": {
              "selector": {"type": "string"},
              "attribute": {"type": "string"}
          }, "required": ["selector"]}),

        T(name="new_tab",
          description="Open a new browser tab.",
          inputSchema={"type": "object", "properties": {
              "url": {"type": "string", "default": "about:blank"}
          }, "required": []}),
        T(name="switch_tab",
          description="Switch to a tab by index number.",
          inputSchema={"type": "object", "properties": {"index": {"type": "integer"}}, "required": ["index"]}),
        T(name="close_tab",
          description="Close the current tab (cannot close the last one).",
          inputSchema={"type": "object", "properties": {}, "required": []}),
        T(name="list_tabs",
          description="List all open tabs with index, URL and title.",
          inputSchema={"type": "object", "properties": {}, "required": []}),

        T(name="get_cookies",   description="Get all cookies for the current context.",
          inputSchema={"type": "object", "properties": {}, "required": []}),
        T(name="set_cookie",
          description="Set a browser cookie.",
          inputSchema={"type": "object", "properties": {
              "name": {"type": "string"}, "value": {"type": "string"},
              "url": {"type": "string"}
          }, "required": ["name", "value"]}),
        T(name="clear_cookies", description="Clear all browser cookies.",
          inputSchema={"type": "object", "properties": {}, "required": []}),
        T(name="clear_cache",
          description="Clear the server-side page content cache (force fresh fetches).",
          inputSchema={"type": "object", "properties": {}, "required": []}),

        T(name="get_console",
          description="Get recent browser console log messages.",
          inputSchema={"type": "object", "properties": {
              "limit": {"type": "integer", "default": 50}
          }, "required": []}),

        T(name="set_fast_mode",
          description="Enable/disable fast mode: blocks images, fonts and CSS to speed up page loads.",
          inputSchema={"type": "object", "properties": {
              "enabled": {"type": "boolean"}
          }, "required": ["enabled"]}),

        # ── Script management ────────────────────────────────────────────
        T(name="script_save",
          description="Save a reusable automation script (Python or JS).",
          inputSchema={"type": "object", "properties": {
              "name": {"type": "string"}, "content": {"type": "string"},
          }, "required": ["name", "content"]}),
        T(name="script_list",   description="List all saved automation scripts.",
          inputSchema={"type": "object", "properties": {}, "required": []}),
        T(name="script_read",
          description="Read source code of a saved script.",
          inputSchema={"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}),
        T(name="script_delete",
          description="Delete a saved script.",
          inputSchema={"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}),
        T(name="script_run",
          description=(
              "Run a saved Python script. Receives current page URL as argv[1]. "
              "Can use playwright. Returns stdout/stderr."
          ),
          inputSchema={"type": "object", "properties": {
              "name": {"type": "string"}, "args": {"type": "string", "default": ""},
              "timeout": {"type": "integer", "default": 30},
          }, "required": ["name"]}),
        T(name="script_run_inline",
          description="Write and run a one-off Python script immediately (not saved).",
          inputSchema={"type": "object", "properties": {
              "code": {"type": "string"}, "timeout": {"type": "integer", "default": 60},
          }, "required": ["code"]}),
    ]


def _cache_key(url: str, kind: str) -> str:
    return f"{kind}:{url}"


@mcp_server.call_tool()
async def call_tool(name: str, arguments: dict):
    try:
        p = await bm.page()

        # ── navigate ────────────────────────────────────────────────────
        if name == "navigate":
            fast = arguments.get("fast_mode", False)
            bm._fast_mode = fast
            resp = await p.goto(
                arguments["url"],
                wait_until=arguments.get("wait_until", "domcontentloaded"),
                timeout=30000,
            )
            status = resp.status if resp else "unknown"
            mode = " [fast: images/css blocked]" if fast else ""
            return ok(f"Navigated → {p.url}  (HTTP {status}){mode}")

        # ── screenshot ──────────────────────────────────────────────────
        if name == "screenshot":
            sel = arguments.get("selector")
            if sel:
                el = await p.query_selector(sel)
                if not el:
                    return ok(f"Element not found: {sel}")
                return img(await el.screenshot())
            return img(await p.screenshot(full_page=arguments.get("full_page", False)))

        # ── get_text (cached) ────────────────────────────────────────────
        if name == "get_text":
            sel = arguments.get("selector", "body")
            url = p.url
            ck = _cache_key(url, f"text:{sel}")
            no_cache = arguments.get("no_cache", False)
            if not no_cache and ck in _page_cache:
                return ok(f"[cached] {_page_cache[ck]}")
            el = await p.query_selector(sel)
            result = await el.inner_text() if el else "Element not found"
            if not no_cache and url != "about:blank":
                _page_cache[ck] = result
            return ok(result)

        # ── get_html (cached) ────────────────────────────────────────────
        if name == "get_html":
            sel = arguments.get("selector", "html")
            outer = arguments.get("outer", True)
            url = p.url
            ck = _cache_key(url, f"html:{sel}:{outer}")
            no_cache = arguments.get("no_cache", False)
            if not no_cache and ck in _page_cache:
                return ok(f"[cached] {_page_cache[ck]}")
            el = await p.query_selector(sel)
            if not el:
                return ok("Element not found")
            result = await el.outer_html() if outer else await el.inner_html()
            if not no_cache and url != "about:blank":
                _page_cache[ck] = result
            return ok(result)

        # ── get_links (cached) ───────────────────────────────────────────
        if name == "get_links":
            limit = arguments.get("limit", 50)
            sel = arguments.get("selector", "a")
            url = p.url
            ck = _cache_key(url, f"links:{sel}:{limit}")
            no_cache = arguments.get("no_cache", False)
            if not no_cache and ck in _page_cache:
                return ok(f"[cached] {_page_cache[ck]}")
            links = await p.evaluate(
                f"""() => Array.from(document.querySelectorAll('{sel}')).slice(0,{limit}).map(a=>{{
                    return {{text:(a.innerText||a.textContent||'').trim().slice(0,100),href:a.href,title:a.title}}
                }})"""
            )
            result = json.dumps(links, ensure_ascii=False)
            if not no_cache and url != "about:blank":
                _page_cache[ck] = result
            return ok(result)

        # ── clear_cache ──────────────────────────────────────────────────
        if name == "clear_cache":
            n = len(_page_cache)
            _page_cache.clear()
            return ok(f"Cache cleared ({n} entries removed)")

        # ── set_fast_mode ────────────────────────────────────────────────
        if name == "set_fast_mode":
            bm._fast_mode = arguments["enabled"]
            state = "enabled" if bm._fast_mode else "disabled"
            return ok(f"Fast mode {state} (images/fonts/CSS {'blocked' if bm._fast_mode else 'allowed'})")

        # ── click ────────────────────────────────────────────────────────
        if name == "click":
            await p.click(arguments["selector"],
                          button=arguments.get("button", "left"),
                          click_count=arguments.get("count", 1))
            return ok(f"Clicked {arguments['selector']}")

        # ── type_text ────────────────────────────────────────────────────
        if name == "type_text":
            if arguments.get("clear"):
                await p.fill(arguments["selector"], "")
            await p.locator(arguments["selector"]).type(arguments["text"])
            return ok(f"Typed into {arguments['selector']}")

        # ── fill ─────────────────────────────────────────────────────────
        if name == "fill":
            await p.fill(arguments["selector"], arguments["value"])
            return ok(f"Filled {arguments['selector']}")

        # ── select_option ─────────────────────────────────────────────────
        if name == "select_option":
            sel = arguments["selector"]
            if v := arguments.get("value"):
                await p.select_option(sel, value=v)
            elif l := arguments.get("label"):
                await p.select_option(sel, label=l)
            else:
                return ok("Provide value or label")
            return ok(f"Selected option in {sel}")

        # ── hover ─────────────────────────────────────────────────────────
        if name == "hover":
            await p.hover(arguments["selector"])
            return ok(f"Hovering over {arguments['selector']}")

        # ── press_key ─────────────────────────────────────────────────────
        if name == "press_key":
            if sel := arguments.get("selector"):
                el = await p.query_selector(sel)
                if el:
                    await el.press(arguments["key"])
                    return ok(f"Pressed {arguments['key']} on {sel}")
            await p.keyboard.press(arguments["key"])
            return ok(f"Pressed {arguments['key']}")

        # ── evaluate ──────────────────────────────────────────────────────
        if name == "evaluate":
            arg_raw = arguments.get("arg")
            arg = json.loads(arg_raw) if arg_raw else None
            result = await p.evaluate(arguments["script"], arg)
            return ok(json.dumps(result, ensure_ascii=False, default=str))

        # ── wait_for ──────────────────────────────────────────────────────
        if name == "wait_for":
            if sel := arguments.get("selector"):
                await p.wait_for_selector(sel, state=arguments.get("state", "visible"), timeout=15000)
                return ok(f"Element {sel} is {arguments.get('state', 'visible')}")
            if ms := arguments.get("ms"):
                await asyncio.sleep(ms / 1000)
                return ok(f"Waited {ms} ms")
            return ok("Provide selector or ms")

        # ── scroll ────────────────────────────────────────────────────────
        if name == "scroll":
            if sel := arguments.get("selector"):
                el = await p.query_selector(sel)
                if el:
                    await el.scroll_into_view_if_needed()
                    return ok(f"Scrolled {sel} into view")
            await p.evaluate(f"window.scrollBy({arguments.get('x', 0)}, {arguments.get('y', 500)})")
            return ok(f"Scrolled by ({arguments.get('x',0)}, {arguments.get('y',500)})")

        if name == "go_back":
            await p.go_back(timeout=10000)
            return ok(f"Back → {p.url}")

        if name == "go_forward":
            await p.go_forward(timeout=10000)
            return ok(f"Forward → {p.url}")

        if name == "reload":
            _page_cache.pop(_cache_key(p.url, "text:body"), None)
            await p.reload()
            return ok(f"Reloaded → {p.url}")

        if name == "get_url":   return ok(p.url)
        if name == "get_title": return ok(await p.title())

        # ── find_elements ─────────────────────────────────────────────────
        if name == "find_elements":
            elements = await p.query_selector_all(arguments["selector"])
            attr = arguments.get("attribute")
            results = []
            for el in elements[:50]:
                results.append(await el.get_attribute(attr) if attr else (await el.inner_text()).strip())
            return ok(json.dumps(results, ensure_ascii=False))

        # ── tabs ──────────────────────────────────────────────────────────
        if name == "new_tab":
            new_page = await bm.new_tab(arguments.get("url", "about:blank"))
            return ok(f"Opened tab {bm.current}: {new_page.url}")

        if name == "switch_tab":
            idx = arguments["index"]
            if 0 <= idx < len(bm.tabs):
                bm.current = idx
                return ok(f"Switched to tab {idx}: {bm.tabs[idx].url}")
            return ok(f"Tab {idx} doesn't exist (have {len(bm.tabs)} tabs)")

        if name == "close_tab":
            if len(bm.tabs) <= 1:
                return ok("Cannot close the last tab")
            await bm.tabs[bm.current].close()
            bm.tabs.pop(bm.current)
            bm.current = max(0, bm.current - 1)
            return ok(f"Closed. Now on tab {bm.current}: {bm.tabs[bm.current].url}")

        if name == "list_tabs":
            result = [{"index": i, "url": t.url, "title": await t.title(), "active": i == bm.current}
                      for i, t in enumerate(bm.tabs)]
            return ok(json.dumps(result, ensure_ascii=False))

        # ── cookies ───────────────────────────────────────────────────────
        if name == "get_cookies":
            return ok(json.dumps(await bm.context.cookies(), ensure_ascii=False, default=str))

        if name == "set_cookie":
            await bm.context.add_cookies([{
                "name": arguments["name"], "value": arguments["value"],
                "url": arguments.get("url", p.url),
            }])
            return ok(f"Set cookie '{arguments['name']}'")

        if name == "clear_cookies":
            await bm.context.clear_cookies()
            return ok("All cookies cleared")

        # ── console ───────────────────────────────────────────────────────
        if name == "get_console":
            limit = arguments.get("limit", 50)
            msgs = bm.console_msgs[-limit:]
            return ok("\n".join(msgs) if msgs else "(no console messages yet)")

        # ── scripts ───────────────────────────────────────────────────────
        if name == "script_save":
            path = SCRIPTS_DIR / arguments["name"]
            path.write_text(arguments["content"])
            return ok(f"Saved: /scripts/{arguments['name']} ({len(arguments['content'])} chars)")

        if name == "script_list":
            files = sorted(SCRIPTS_DIR.iterdir())
            if not files:
                return ok("No scripts saved yet.")
            return ok("\n".join(f"{f.name}  ({f.stat().st_size} bytes)" for f in files))

        if name == "script_read":
            path = SCRIPTS_DIR / arguments["name"]
            return ok(path.read_text() if path.exists() else f"Script not found: {arguments['name']}")

        if name == "script_delete":
            path = SCRIPTS_DIR / arguments["name"]
            if not path.exists():
                return ok(f"Script not found: {arguments['name']}")
            path.unlink()
            return ok(f"Deleted: {arguments['name']}")

        if name == "script_run":
            path = SCRIPTS_DIR / arguments["name"]
            if not path.exists():
                return ok(f"Script not found: {arguments['name']}")
            args_str = arguments.get("args", "")
            timeout = min(int(arguments.get("timeout", 30)), 300)
            cmd = ["python3", str(path), p.url] + (args_str.split() if args_str else [])
            try:
                r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
                out = []
                if r.stdout: out.append("STDOUT:\n" + r.stdout.rstrip())
                if r.stderr: out.append("STDERR:\n" + r.stderr.rstrip())
                out.append(f"Exit: {r.returncode}")
                return ok("\n\n".join(out))
            except subprocess.TimeoutExpired:
                return ok("[TIMEOUT]")

        if name == "script_run_inline":
            timeout = min(int(arguments.get("timeout", 60)), 300)
            tmp = Path("/tmp/_inline_script.py")
            tmp.write_text(arguments["code"])
            try:
                r = subprocess.run(["python3", str(tmp)], capture_output=True, text=True, timeout=timeout)
                out = []
                if r.stdout: out.append("STDOUT:\n" + r.stdout.rstrip())
                if r.stderr: out.append("STDERR:\n" + r.stderr.rstrip())
                out.append(f"Exit: {r.returncode}")
                return ok("\n\n".join(out))
            except subprocess.TimeoutExpired:
                return ok("[TIMEOUT]")

        return ok(f"Unknown tool: {name}")

    except Exception as e:
        return ok(f"[ERROR] {type(e).__name__}: {e}")


async def handle_sse(request: Request):
    async with sse.connect_sse(request.scope, request.receive, request._send) as streams:
        await mcp_server.run(streams[0], streams[1], mcp_server.create_initialization_options())


async def handle_health(request: Request):
    cached = len(_page_cache)
    fast = "on" if bm._fast_mode else "off"
    status = f"browser-mcp v2 OK — tabs:{len(bm.tabs)} current:{bm.current} cache:{cached} fast_mode:{fast}"
    return Response(status, media_type="text/plain")


app = Starlette(routes=[
    Route("/sse", endpoint=handle_sse),
    Route("/health", endpoint=handle_health),
    Mount("/messages/", app=sse.handle_post_message),
])

if __name__ == "__main__":
    threading.Thread(target=lambda: asyncio.run(bm.init()), daemon=True).start()
    uvicorn.run(app, host="0.0.0.0", port=7860)