"""web_view — headless browser page extraction."""

from __future__ import annotations

import asyncio
import base64
from contextlib import suppress
from typing import Any

from playwright.async_api import async_playwright

from .extract import clean_visible_text, extract_page, truncate

_TIMEOUT = 20_000
_MAX_TEXT = 20_000

# A call may open several pages at once: pages take 1-8s each, so a batch costs
# the slowest instead of the sum (three URLs here: 17.3s one by one, 1.9s
# together). The shared text budget is the counterweight — three pages at the
# single-page cap would be 60k chars of context, so a batch stops at 40k and
# says where it stopped.
_MAX_BATCH = 3
_BATCH_TEXT = 40_000
_BATCH_DEADLINE = 30.0

# One Chromium for the life of the server. Launching a browser per call cost ~5s
# even for ``example.com``; this process is long-lived, so the browser can be
# too. Each call still gets a fresh context, which is what isolates cookies,
# storage and the page from its neighbours.
_playwright = None
_browser = None
_browser_lock = asyncio.Lock()

# Contexts are cheap but not free; the cap keeps a burst of concurrent calls
# (a batch, or two tool calls Navi ran in parallel) from launching a crowd.
_view_slots = asyncio.Semaphore(_MAX_BATCH)


def _normalize(raw: str | list[str]) -> tuple[list[str], list[str]]:
    """Turn the argument into the URLs to open, plus notes worth telling the model."""
    urls: list[str] = []
    for item in [raw] if isinstance(raw, str) else raw:
        url = str(item).strip()
        if url and url not in urls:
            urls.append(url)
    notes: list[str] = []
    if len(urls) > _MAX_BATCH:
        notes.append(
            f"Only the first {_MAX_BATCH} of {len(urls)} URLs were opened — "
            "ask the rest in another call."
        )
        urls = urls[:_MAX_BATCH]
    return urls, notes



async def _get_browser():
    """Return the shared browser, launching (or relaunching) it if needed."""
    global _playwright, _browser
    async with _browser_lock:
        if _browser is not None and _browser.is_connected():
            return _browser
        # A browser that crashed leaves a Playwright driver behind; stop it
        # before starting another, or the driver processes accumulate.
        if _playwright is not None:
            with suppress(Exception):
                await _playwright.stop()
            _playwright = None
            _browser = None
        _playwright = await async_playwright().start()
        _browser = await _playwright.chromium.launch(headless=True)
        return _browser


async def close_browser() -> None:
    """Stop Chromium and its Playwright driver. Called from the server lifespan."""
    global _playwright, _browser
    async with _browser_lock:
        if _browser is not None:
            with suppress(Exception):
                await _browser.close()
            _browser = None
        if _playwright is not None:
            with suppress(Exception):
                await _playwright.stop()
            _playwright = None


# The page's own visible text, without the chrome that never had a reader.
_VISIBLE_TEXT = """() => {
    const kill = ['script','style','noscript','iframe',
                  'nav','header','footer','aside',
                  '[role="navigation"]','[role="banner"]',
                  '[role="contentinfo"]'];
    const clone = document.body.cloneNode(true);
    kill.forEach(sel => {
        clone.querySelectorAll(sel).forEach(el => el.remove());
    });
    return clone.innerText || clone.textContent || '';
}"""


async def web_view(
    url: str | list[str],
    screenshot: bool = False,
    wait_until: str = "networkidle",
) -> dict[str, Any]:
    """Open a URL — or up to three — in a headless browser and return readable text.

    Returns dict with success, output, error, metadata.
    """
    urls, notes = _normalize(url)
    if not urls:
        return {
            "success": False,
            "output": "No URL given.",
            "error": "empty_url",
            "metadata": {},
        }
    if screenshot and len(urls) > 1:
        return {
            "success": False,
            "output": (
                "screenshot works on one URL at a time — ask for a single URL "
                "with screenshot, or several without."
            ),
            "error": "screenshot_batch",
            "metadata": {},
        }

    if len(urls) == 1 and not notes:
        return await _view_bounded(urls[0], screenshot=screenshot, wait_until=wait_until)

    outcomes = await asyncio.gather(
        *(_view_bounded(u, wait_until=wait_until) for u in urls),
        return_exceptions=True,
    )
    return _combine(urls, outcomes, notes)


async def _view_bounded(
    url: str,
    screenshot: bool = False,
    wait_until: str = "networkidle",
) -> dict[str, Any]:
    """One page, bounded on its own.

    The deadline is per URL: a batch waits for its slowest page, and a single
    ceiling around the whole batch would let one hung page erase the others
    that already loaded.
    """
    try:
        return await asyncio.wait_for(
            _view_one(url, screenshot=screenshot, wait_until=wait_until),
            timeout=_BATCH_DEADLINE,
        )
    except TimeoutError:
        return {
            "success": False,
            "output": f"No answer within {_BATCH_DEADLINE:.0f}s: {url}",
            "error": "timeout",
            "metadata": {},
        }


def _combine(
    urls: list[str],
    outcomes: list[Any],
    notes: list[str],
) -> dict[str, Any]:
    """One block per URL, sharing one text budget.

    A URL that failed is named rather than dropped — the model can retry that
    one page instead of repeating the batch — and a batch that opened something
    is a success even when the rest did not.
    """
    blocks: list[str] = []
    per_url: list[dict[str, Any]] = []
    budget = _BATCH_TEXT
    opened = 0

    for url, outcome in zip(urls, outcomes, strict=True):
        if isinstance(outcome, BaseException):
            blocks.append(f"### {url} -> failed: {outcome}")
            per_url.append({"url": url, "success": False, "error": str(outcome)})
            continue
        if not outcome["success"]:
            blocks.append(f"### {url} -> failed: {outcome['output']}")
            per_url.append({"url": url, "success": False, "error": outcome["error"]})
            continue

        opened += 1
        block = f"### {url}\n{outcome['output']}"
        if budget <= 0:
            block = (
                f"### {url}\n[… not shown: the batch's "
                f"{_BATCH_TEXT // 1000}k-char text budget is used up]"
            )
            per_url.append({"url": url, "success": True, "truncated": True})
            blocks.append(block)
            continue
        if len(block) > budget:
            block = (
                block[:budget]
                + f"\n\n[… truncated: the batch's {_BATCH_TEXT // 1000}k-char "
                "text budget is used up]"
            )
            per_url.append({"url": url, "success": True, "truncated": True})
            budget = 0
        else:
            per_url.append({"url": url, "success": True})
            budget -= len(block)
        blocks.append(block)

    body = "\n\n".join(notes + blocks)
    if not opened:
        return {
            "success": False,
            "output": body or "No URL could be opened.",
            "error": "all_urls_failed",
            "metadata": {"urls": per_url},
        }
    return {
        "success": True,
        "output": body,
        "error": None,
        "metadata": {"urls": per_url},
    }


async def _view_one(
    url: str,
    screenshot: bool = False,
    wait_until: str = "networkidle",
) -> dict[str, Any]:
    """Open one page in its own context. Returns dict with success, output, error, metadata."""
    if not url.startswith(("http://", "https://")):
        return {
            "success": False,
            "output": "URL must start with http:// or https://",
            "error": "invalid_url",
            "metadata": {},
        }

    try:
        async with _view_slots:
            browser = await _get_browser()
            context = await browser.new_context(
                viewport={"width": 1280, "height": 800},
                user_agent=(
                    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
                    "(KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
                ),
            )
            try:
                return await _read_page(context, url, screenshot, wait_until)
            finally:
                with suppress(Exception):
                    await context.close()

    except Exception as e:
        return {
            "success": False,
            "output": f"Browser error: {e}",
            "error": str(e),
            "metadata": {},
        }


async def _read_page(
    context: Any,
    url: str,
    screenshot: bool,
    wait_until: str,
) -> dict[str, Any]:
    page = await context.new_page()

    try:
        await page.goto(url, wait_until=wait_until, timeout=_TIMEOUT)
    except Exception:
        pass

    title = await page.title()
    final_url = page.url

    # The visible text is both the fallback and the yardstick: a page whose
    # extractor kept a fraction of it is a page the extractor did not understand.
    visible = clean_visible_text(await page.evaluate(_VISIBLE_TEXT))
    extracted = extract_page(await page.content(), final_url, visible_chars=len(visible))

    if extracted is not None:
        text, links = extracted
        text = truncate(text, _MAX_TEXT)
        if links:
            text = f"{text}\n\nLinks on this page:\n" + "\n".join(links)
        extractor = "trafilatura"
    else:
        text, links, extractor = truncate(visible, _MAX_TEXT), [], "innerText"

    output_parts = []
    if title:
        output_parts.append(f"Title: {title}")
    if final_url != url:
        output_parts.append(f"Final URL: {final_url}")
    output_parts.append("")
    output_parts.append(text)
    output = "\n".join(output_parts)

    metadata: dict[str, Any] = {"extractor": extractor, "links": len(links)}
    if screenshot:
        png = await page.screenshot(full_page=False)
        metadata |= {
            "base64": base64.b64encode(png).decode(),
            "mime": "image/png",
            "is_image": True,
        }

    return {
        "success": True,
        "output": output,
        "error": None,
        "metadata": metadata,
    }