core-python/maris_core/browser/automation.py

"""Playwright-powered browser automation endpoints for Maris AI."""

from __future__ import annotations

import base64
import logging
import uuid
from typing import Any, Literal, NoReturn, Protocol
from urllib.parse import urlsplit

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field, field_validator

logger = logging.getLogger(__name__)
router = APIRouter()

BrowserWaitUntil = Literal["load", "domcontentloaded", "networkidle", "commit"]
ALLOWED_BROWSER_URL_SCHEMES = ("about", "data", "http", "https")
# Conservative cap keeps Playwright resource usage bounded on shared runtimes.
MAX_BROWSER_AUTOMATION_SESSIONS = 4


class BrowserAutomationUnavailableError(RuntimeError):
    """Raised when the Playwright runtime is not installed or ready."""


class BrowserAutomationSession(Protocol):
    """Protocol for a browser session backend."""

    headless: bool
    viewport: dict[str, int]

    async def snapshot(self) -> dict[str, Any]: ...

    async def navigate(
        self, url: str, *, wait_until: BrowserWaitUntil, timeout_ms: int
    ) -> dict[str, Any]: ...

    async def click(
        self,
        selector: str,
        *,
        timeout_ms: int,
        wait_until_after: BrowserWaitUntil | None = None,
    ) -> dict[str, Any]: ...

    async def fill(
        self,
        selector: str,
        value: str,
        *,
        timeout_ms: int,
        submit: bool,
    ) -> dict[str, Any]: ...

    async def extract_text(
        self, selector: str | None, *, timeout_ms: int, max_length: int
    ) -> str: ...

    async def screenshot_png(self, *, full_page: bool) -> bytes: ...

    async def close(self) -> None: ...


class BrowserSessionStartRequest(BaseModel):
    headless: bool = True
    viewport_width: int = Field(default=1440, ge=320, le=3840)
    viewport_height: int = Field(default=900, ge=240, le=2160)


class BrowserSessionRequest(BaseModel):
    session_id: str = Field(min_length=1, max_length=120)


class BrowserNavigateRequest(BrowserSessionRequest):
    url: str = Field(min_length=1, max_length=4096)
    wait_until: BrowserWaitUntil = "domcontentloaded"
    timeout_ms: int = Field(default=15000, ge=1000, le=120000)

    @field_validator("url")
    @classmethod
    def validate_url(cls, value: str) -> str:
        return _normalize_browser_url(value)


class BrowserClickRequest(BrowserSessionRequest):
    selector: str = Field(min_length=1, max_length=1024)
    timeout_ms: int = Field(default=8000, ge=500, le=120000)
    wait_until_after: BrowserWaitUntil | None = None


class BrowserFillRequest(BrowserSessionRequest):
    selector: str = Field(min_length=1, max_length=1024)
    value: str = Field(max_length=4000)
    timeout_ms: int = Field(default=8000, ge=500, le=120000)
    submit: bool = False


class BrowserExtractRequest(BrowserSessionRequest):
    selector: str | None = Field(default=None, max_length=1024)
    timeout_ms: int = Field(default=8000, ge=500, le=120000)
    max_length: int = Field(default=4000, ge=1, le=12000)


class BrowserScreenshotRequest(BrowserSessionRequest):
    full_page: bool = False


class BrowserSessionResponse(BaseModel):
    session_id: str
    active: bool
    headless: bool
    viewport: dict[str, int]
    url: str | None = None
    title: str | None = None


class BrowserTextResponse(BaseModel):
    session_id: str
    url: str | None
    title: str | None
    text: str


class BrowserScreenshotResponse(BaseModel):
    session_id: str
    url: str | None
    title: str | None
    image_base64: str
    mime_type: str = "image/png"


class BrowserCapabilitiesResponse(BaseModel):
    provider: str
    package: str
    install_command: str
    browser_install_command: str
    supported_actions: list[str]
    allowed_url_schemes: list[str]
    max_sessions: int
    headless_default: bool


class ManagedPlaywrightSession:
    """Thin wrapper around a Playwright Chromium page."""

    def __init__(
        self, *, playwright: Any, browser: Any, page: Any, headless: bool, viewport: dict[str, int]
    ):
        self._playwright = playwright
        self._browser = browser
        self._page = page
        self.headless = headless
        self.viewport = viewport

    @classmethod
    async def create(cls, *, headless: bool, viewport: dict[str, int]) -> ManagedPlaywrightSession:
        try:
            from playwright.async_api import async_playwright
        except ImportError as exc:
            raise BrowserAutomationUnavailableError(
                "Browser automation vajag Playwright. Instalē ar `pip install playwright` "
                "un `python -m playwright install chromium`."
            ) from exc

        playwright = await async_playwright().start()
        try:
            browser = await playwright.chromium.launch(headless=headless)
            page = await browser.new_page(viewport=viewport)
        except _browser_operation_errors():
            await playwright.stop()
            raise
        return cls(
            playwright=playwright,
            browser=browser,
            page=page,
            headless=headless,
            viewport=viewport,
        )

    async def snapshot(self) -> dict[str, Any]:
        title = None
        try:
            title = await self._page.title()
        except _browser_operation_errors():
            logger.debug("Browser page title unavailable", exc_info=True)
        return {"url": self._page.url or None, "title": title}

    async def navigate(
        self, url: str, *, wait_until: BrowserWaitUntil, timeout_ms: int
    ) -> dict[str, Any]:
        await self._page.goto(url, wait_until=wait_until, timeout=timeout_ms)
        return await self.snapshot()

    async def click(
        self,
        selector: str,
        *,
        timeout_ms: int,
        wait_until_after: BrowserWaitUntil | None = None,
    ) -> dict[str, Any]:
        locator = self._page.locator(selector).first
        await locator.wait_for(state="visible", timeout=timeout_ms)
        await locator.click(timeout=timeout_ms)
        if wait_until_after is not None:
            await self._page.wait_for_load_state(wait_until_after, timeout=timeout_ms)
        return await self.snapshot()

    async def fill(
        self,
        selector: str,
        value: str,
        *,
        timeout_ms: int,
        submit: bool,
    ) -> dict[str, Any]:
        locator = self._page.locator(selector).first
        await locator.wait_for(state="visible", timeout=timeout_ms)
        await locator.fill(value, timeout=timeout_ms)
        if submit:
            await locator.press("Enter", timeout=timeout_ms)
        return await self.snapshot()

    async def extract_text(self, selector: str | None, *, timeout_ms: int, max_length: int) -> str:
        locator = self._page.locator(selector or "body").first
        await locator.wait_for(state="attached", timeout=timeout_ms)
        text = (await locator.inner_text(timeout=timeout_ms)).strip()
        return text[:max_length]

    async def screenshot_png(self, *, full_page: bool) -> bytes:
        return await self._page.screenshot(type="png", full_page=full_page)

    async def close(self) -> None:
        await self._browser.close()
        await self._playwright.stop()


_browser_sessions: dict[str, BrowserAutomationSession] = {}


def get_browser_automation_capabilities() -> BrowserCapabilitiesResponse:
    """Return public browser automation capability metadata."""
    return BrowserCapabilitiesResponse(
        provider="playwright",
        package="playwright",
        install_command="pip install playwright",
        browser_install_command="python -m playwright install chromium",
        supported_actions=[
            "navigate",
            "click",
            "fill",
            "extract_text",
            "screenshot",
            "close_session",
        ],
        allowed_url_schemes=list(ALLOWED_BROWSER_URL_SCHEMES),
        max_sessions=MAX_BROWSER_AUTOMATION_SESSIONS,
        headless_default=True,
    )


def _normalize_browser_url(url: str) -> str:
    normalized = url.strip()
    parsed = urlsplit(normalized)
    if parsed.scheme not in ALLOWED_BROWSER_URL_SCHEMES:
        allowed_schemes = ", ".join(f"{scheme}:" for scheme in ALLOWED_BROWSER_URL_SCHEMES)
        raise ValueError(f"Browser automation allows only these URL schemes: {allowed_schemes}.")
    if parsed.scheme in {"http", "https"} and not parsed.netloc:
        raise ValueError("HTTP/HTTPS URL must include a valid hostname.")
    return normalized


def _browser_operation_errors() -> tuple[type[BaseException], ...]:
    """Return browser-operation exceptions for runtimes with or without Playwright installed."""
    errors: list[type[BaseException]] = [
        BrowserAutomationUnavailableError,
        RuntimeError,
        TimeoutError,
        ValueError,
    ]
    try:
        from playwright.async_api import Error as PlaywrightError
        from playwright.async_api import TimeoutError as PlaywrightTimeoutError
    except ImportError:
        return tuple(errors)
    return tuple([*errors, PlaywrightError, PlaywrightTimeoutError])


async def _create_browser_session(
    *, headless: bool, viewport: dict[str, int]
) -> BrowserAutomationSession:
    return await ManagedPlaywrightSession.create(headless=headless, viewport=viewport)


def _get_browser_session(session_id: str) -> BrowserAutomationSession:
    session = _browser_sessions.get(session_id)
    if session is None:
        raise HTTPException(status_code=404, detail="Browser automation sesija netika atrasta.")
    return session


def _snapshot_response(
    session_id: str, session: BrowserAutomationSession, snapshot: dict[str, Any]
) -> BrowserSessionResponse:
    return BrowserSessionResponse(
        session_id=session_id,
        active=True,
        headless=session.headless,
        viewport=session.viewport,
        url=snapshot.get("url"),
        title=snapshot.get("title"),
    )


def _raise_browser_http_error(exc: Exception) -> NoReturn:
    if isinstance(exc, HTTPException):
        raise exc
    if isinstance(exc, BrowserAutomationUnavailableError):
        raise HTTPException(status_code=503, detail=str(exc)) from exc
    if isinstance(exc, ValueError):
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    message = str(exc).strip() or "Browser automation darbība neizdevās."
    lower = message.lower()
    if "timeout" in lower:
        raise HTTPException(status_code=504, detail=message) from exc
    raise HTTPException(status_code=502, detail=message) from exc


@router.get("/capabilities", response_model=BrowserCapabilitiesResponse)
async def browser_capabilities() -> BrowserCapabilitiesResponse:
    """Return the browser automation surface available in Maris."""
    return get_browser_automation_capabilities()


@router.post("/sessions/start", response_model=BrowserSessionResponse)
async def start_browser_session(req: BrowserSessionStartRequest) -> BrowserSessionResponse:
    """Create a safe, isolated browser automation session."""
    if len(_browser_sessions) >= MAX_BROWSER_AUTOMATION_SESSIONS:
        raise HTTPException(
            status_code=429,
            detail="Sasniegts maksimālais browser automation sesiju limits.",
        )

    viewport = {"width": req.viewport_width, "height": req.viewport_height}
    session_id = str(uuid.uuid4())

    try:
        session = await _create_browser_session(headless=req.headless, viewport=viewport)
        _browser_sessions[session_id] = session
        snapshot = await session.snapshot()
        return _snapshot_response(session_id, session, snapshot)
    except _browser_operation_errors() as exc:
        _browser_sessions.pop(session_id, None)
        _raise_browser_http_error(exc)


@router.post("/navigate", response_model=BrowserSessionResponse)
async def navigate_browser(req: BrowserNavigateRequest) -> BrowserSessionResponse:
    """Navigate an existing browser automation session to a URL."""
    session = _get_browser_session(req.session_id)
    try:
        snapshot = await session.navigate(
            req.url, wait_until=req.wait_until, timeout_ms=req.timeout_ms
        )
        return _snapshot_response(req.session_id, session, snapshot)
    except _browser_operation_errors() as exc:
        _raise_browser_http_error(exc)


@router.post("/click", response_model=BrowserSessionResponse)
async def click_browser(req: BrowserClickRequest) -> BrowserSessionResponse:
    """Click an element in the current browser session."""
    session = _get_browser_session(req.session_id)
    try:
        snapshot = await session.click(
            req.selector,
            timeout_ms=req.timeout_ms,
            wait_until_after=req.wait_until_after,
        )
        return _snapshot_response(req.session_id, session, snapshot)
    except _browser_operation_errors() as exc:
        _raise_browser_http_error(exc)


@router.post("/fill", response_model=BrowserSessionResponse)
async def fill_browser(req: BrowserFillRequest) -> BrowserSessionResponse:
    """Fill an input in the current browser session."""
    session = _get_browser_session(req.session_id)
    try:
        snapshot = await session.fill(
            req.selector,
            req.value,
            timeout_ms=req.timeout_ms,
            submit=req.submit,
        )
        return _snapshot_response(req.session_id, session, snapshot)
    except _browser_operation_errors() as exc:
        _raise_browser_http_error(exc)


@router.post("/extract", response_model=BrowserTextResponse)
async def extract_browser_text(req: BrowserExtractRequest) -> BrowserTextResponse:
    """Extract visible text from the page or a specific selector."""
    session = _get_browser_session(req.session_id)
    try:
        snapshot = await session.snapshot()
        text = await session.extract_text(
            req.selector, timeout_ms=req.timeout_ms, max_length=req.max_length
        )
        return BrowserTextResponse(
            session_id=req.session_id,
            url=snapshot.get("url"),
            title=snapshot.get("title"),
            text=text,
        )
    except _browser_operation_errors() as exc:
        _raise_browser_http_error(exc)


@router.post("/screenshot", response_model=BrowserScreenshotResponse)
async def screenshot_browser(req: BrowserScreenshotRequest) -> BrowserScreenshotResponse:
    """Capture a PNG screenshot from the current browser session."""
    session = _get_browser_session(req.session_id)
    try:
        snapshot = await session.snapshot()
        png = await session.screenshot_png(full_page=req.full_page)
        return BrowserScreenshotResponse(
            session_id=req.session_id,
            url=snapshot.get("url"),
            title=snapshot.get("title"),
            image_base64=base64.b64encode(png).decode("ascii"),
        )
    except _browser_operation_errors() as exc:
        _raise_browser_http_error(exc)


@router.post("/sessions/close", response_model=BrowserSessionResponse)
async def close_browser_session(req: BrowserSessionRequest) -> BrowserSessionResponse:
    """Close a browser automation session and release resources."""
    session = _browser_sessions.pop(req.session_id, None)
    if session is None:
        raise HTTPException(status_code=404, detail="Browser automation sesija netika atrasta.")

    try:
        snapshot = await session.snapshot()
    except _browser_operation_errors():
        snapshot = {"url": None, "title": None}

    try:
        await session.close()
    except _browser_operation_errors() as exc:
        _raise_browser_http_error(exc)

    return BrowserSessionResponse(
        session_id=req.session_id,
        active=False,
        headless=session.headless,
        viewport=session.viewport,
        url=snapshot.get("url"),
        title=snapshot.get("title"),
    )