miniswardbower/browser/controller.py

"""
High-Performance Playwright Browser Controller.
Manages browser lifecycle, isolated contexts, stealth flags, fast execution,
and visual annotated screenshot auditing.
"""

from __future__ import annotations

import asyncio
import io
import logging
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union
from PIL import Image, ImageDraw
from playwright.async_api import Browser, BrowserContext, Page, Error as PlaywrightError, async_playwright

from miniswardbower.browser.action_executor import ActionExecutor
from miniswardbower.browser.dom_distiller import DOMDistiller
from miniswardbower.core.config import BrowserConfig
from miniswardbower.core.schemas import ActionBatch, BrowserAction, BrowserActionType, PrunedAXTree


class BrowserController:
    """Async Playwright browser engine for the multi-agent team."""

    def __init__(self, config: Optional[BrowserConfig] = None):
        self.config = config or BrowserConfig()
        self._playwright = None
        self._browser: Optional[Browser] = None
        self._context: Optional[BrowserContext] = None
        self._page: Optional[Page] = None
        self._lightpanda_driver = None
        self._obscura_driver = None
        self.distiller = DOMDistiller()
        # Obscura 0.2.2 hit-testing lets pointer-events:none overlays intercept
        # clicks. Keep annotations in the captured image until upstream fixes it.
        if self.config.browser_engine == "obscura":
            self.distiller.draw_overlays = False
        self.executor = ActionExecutor(action_timeout_ms=self.config.action_timeout_ms, enable_natural_dynamics=True)

    async def start(self) -> None:
        """Launches the configured browser; failures never silently change engines."""
        if self._playwright is not None or self._lightpanda_driver is not None or self._obscura_driver is not None:
            return

        if self.config.browser_engine in ("lightpanda", "obscura"):
            if self.config.browser_engine == "obscura":
                from miniswardbower.browser.obscura_driver import ObscuraDriver
                driver = ObscuraDriver(self.config.obscura_bin_path, self.config.obscura_allow_private_network)
            else:
                from miniswardbower.browser.lightpanda_driver import LightpandaDriver
                driver = LightpandaDriver(bin_path=self.config.lightpanda_bin_path)
            try:
                self._browser, self._context, self._page = await driver.start()
                await self.page.set_viewport_size({"width": self.config.window_width, "height": self.config.window_height})
            except BaseException:
                await driver.stop()
                self._browser = self._context = self._page = None
                raise
            setattr(self, f"_{self.config.browser_engine}_driver", driver)
            return

        self._playwright = await async_playwright().start()
        self._browser = await self._playwright.chromium.launch(
            headless=self.config.headless,
            args=[
                "--disable-blink-features=AutomationControlled",
                "--no-sandbox",
                "--disable-setuid-sandbox",
                "--disable-infobars",
                "--window-size=1280,800",
            ],
        )
        self._context = await self._browser.new_context(
            viewport={"width": self.config.window_width, "height": self.config.window_height},
            user_agent=self.config.user_agent,
            device_scale_factor=1,
            accept_downloads=True,
            locale="en-US",
            timezone_id="America/New_York",
            permissions=["geolocation", "notifications"],
            extra_http_headers={
                "Accept-Language": "en-US,en;q=0.9",
                "Sec-Ch-Ua": '"Not A(Brand";v="99", "Google Chrome";v="121", "Chromium";v="121"',
                "Sec-Ch-Ua-Mobile": "?0",
                "Sec-Ch-Ua-Platform": '"Linux"',
            },
        )

        # Stealth CDP Profile: Mask automation signatures & emulate authentic physical browser
        await self._context.add_init_script("""
            // 1. Mask navigator.webdriver
            try {
                Object.defineProperty(navigator, 'webdriver', { get: () => undefined });
            } catch (e) {}

            // 2. Realistic window.chrome object
            if (!window.chrome) {
                window.chrome = {
                    app: { isInstalled: false },
                    runtime: { OnInstalledReason: {}, OnRestartRequiredReason: {} }
                };
            }

            // 3. Realistic plugins list
            try {
                Object.defineProperty(navigator, 'plugins', {
                    get: () => [
                        { name: 'Chrome PDF Plugin', filename: 'internal-pdf-viewer', description: 'Portable Document Format' },
                        { name: 'Chrome PDF Viewer', filename: 'mhjfbmdgcfjbbpaeojofohoefgiehjai', description: '' },
                        { name: 'Native Client', filename: 'internal-nacl-plugin', description: '' }
                    ]
                });
            } catch (e) {}

            // 4. Genuine languages
            try {
                Object.defineProperty(navigator, 'languages', { get: () => ['en-US', 'en'] });
            } catch (e) {}

            // 5. Hardware-accelerated WebGL vendor profile (NVIDIA)
            try {
                const origGetParam = WebGLRenderingContext.prototype.getParameter;
                WebGLRenderingContext.prototype.getParameter = function(param) {
                    if (param === 37445) return 'Google Inc. (NVIDIA)';
                    if (param === 37446) return 'ANGLE (NVIDIA, NVIDIA GeForce RTX 5090 Direct3D11 vs_5_0 ps_5_0, D3D11)';
                    return origGetParam.apply(this, arguments);
                };
            } catch (e) {}
        """)

        self._page = await self._context.new_page()

    async def stop(self) -> None:
        """Closes browser context and its driver."""
        if self._obscura_driver is not None:
            await self._obscura_driver.stop()
            self._obscura_driver = None
            self._page = self._context = self._browser = None
            return

        if self._lightpanda_driver is not None:
            await self._lightpanda_driver.stop()
            self._lightpanda_driver = None
            self._page = None
            self._context = None
            self._browser = None
            return

        # A stalled page must not prevent closing its browser and transport.
        for resource, method in ((self._page, "close"), (self._context, "close"),
                                 (self._browser, "close"), (self._playwright, "stop")):
            if resource is not None:
                try:
                    await asyncio.wait_for(getattr(resource, method)(), timeout=5)
                except Exception as exc:
                    logging.getLogger(__name__).warning("Browser cleanup %s failed: %s", method, exc)
        self._page = None
        self._context = None
        self._browser = None
        self._playwright = None

    @property
    def page(self) -> Page:
        if self._page is None:
            raise RuntimeError("Browser is not started. Call await controller.start() first.")
        return self._page

    async def goto(self, url: str) -> None:
        """Navigates to URL and waits for DOM content loaded."""
        await self.page.goto(url, wait_until="domcontentloaded", timeout=self.config.navigation_timeout_ms)

    async def get_pruned_tree(self) -> PrunedAXTree:
        """Extracts the lightweight, sanitized AXTree."""
        for attempt in range(3):
            try:
                return await self.distiller.extract_pruned_tree(self.page)
            except PlaywrightError as exc:
                if "Execution context was destroyed" not in str(exc) or attempt == 2:
                    raise
                await self.page.wait_for_load_state("domcontentloaded", timeout=self.config.navigation_timeout_ms)
                await asyncio.sleep(0.2)

    async def screenshot(self, output_path: Union[Path, str]) -> Path:
        """Captures viewport screenshot for vision grounding or audit trail."""
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        await self.page.screenshot(path=str(output_path), full_page=False)
        return output_path

    async def capture_annotated_step_screenshot(
        self,
        output_path: Union[Path, str],
        step_num: int,
        action_label: str,
        bbox: Optional[Tuple[float, float, float, float]] = None,
        click_coords: Optional[Tuple[float, float]] = None,
    ) -> Path:
        """
        Captures a high-resolution screenshot with visual annotations:
        - Target element bounding box with translucent green tint
        - Click coordinate glowing crosshairs
        - Step execution HUD badge
        """
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        raw_bytes = await self.page.screenshot(full_page=False)
        img = Image.open(io.BytesIO(raw_bytes)).convert("RGBA")

        overlay = Image.new("RGBA", img.size, (0, 0, 0, 0))
        draw = ImageDraw.Draw(overlay)

        # 1. Bounding box overlay
        if bbox and bbox[2] > 0 and bbox[3] > 0:
            bx, by, bw, bh = bbox
            draw.rectangle(
                [bx, by, bx + bw, by + bh],
                fill=(16, 185, 129, 45),
                outline=(16, 185, 129, 255),
                width=3,
            )

        # 2. Click coordinate crosshair overlay
        if click_coords:
            cx, cy = click_coords
            # Outer ring
            draw.ellipse([cx - 14, cy - 14, cx + 14, cy + 14], outline=(239, 68, 68, 220), width=3)
            # Center target
            draw.ellipse([cx - 4, cy - 4, cx + 4, cy + 4], fill=(239, 68, 68, 255))
            # Crosshair tick marks
            draw.line([cx - 20, cy, cx - 6, cy], fill=(239, 68, 68, 240), width=2)
            draw.line([cx + 6, cy, cx + 20, cy], fill=(239, 68, 68, 240), width=2)
            draw.line([cx, cy - 20, cx, cy - 6], fill=(239, 68, 68, 240), width=2)
            draw.line([cx, cy + 6, cx, cy + 20], fill=(239, 68, 68, 240), width=2)

        # 3. Top-left HUD badge
        badge_width = min(420, max(280, len(action_label) * 8 + 60))
        draw.rectangle([12, 12, 12 + badge_width, 50], fill=(15, 23, 42, 225), outline=(56, 189, 248, 255), width=2)
        draw.text((22, 20), f"Step {step_num}: {action_label[:45]}", fill=(248, 250, 252, 255))

        annotated = Image.alpha_composite(img, overlay).convert("RGB")
        annotated.save(output_path, "PNG")
        return output_path

    async def execute_action(self, action: BrowserAction, tree: Optional[PrunedAXTree] = None) -> Dict[str, Any]:
        """Executes a single atomic action."""
        return await self.executor.execute_action(self.page, action, tree)

    async def execute_batch(self, batch: ActionBatch, tree: Optional[PrunedAXTree] = None) -> List[Dict[str, Any]]:
        """Executes a speculative batch of actions."""
        return await self.executor.execute_batch(self.page, batch, tree)

    async def execute_chunk(
        self,
        actions: List[BrowserAction],
        tree: Optional[PrunedAXTree] = None,
        early_stop_on_navigation: bool = True,
    ) -> List[Dict[str, Any]]:
        """Executes a pipelined speculative action chunk."""
        return await self.executor.execute_chunk(
            self.page, actions, tree, early_stop_on_navigation=early_stop_on_navigation
        )

    async def scrape_page(self) -> Dict[str, Any]:
        """Scrapes structured page data (metadata, tables, card elements, text)."""
        action = BrowserAction(op=BrowserActionType.EXTRACT)
        details = await self.executor.execute_action(self.page, action)
        return details.get("extracted_data", {})

    async def scrape_element(self, target: str, tree: Optional[PrunedAXTree] = None) -> Dict[str, Any]:
        """Scrapes structured content and attributes from a specific targeted element."""
        action = BrowserAction(op=BrowserActionType.EXTRACT, target=target)
        details = await self.executor.execute_action(self.page, action, tree)
        return details.get("extracted_data", {})

    async def scrape_table(self, selector: str = "table") -> List[List[str]]:
        """Scrapes 2D matrix rows from a tabular element."""
        action = BrowserAction(op=BrowserActionType.EXTRACT, extract_selector=selector)
        details = await self.executor.execute_action(self.page, action)
        data = details.get("extracted_data", [])
        return data