"""A real Chromium for the web agent, driven with Playwright.

One Chromium process serves every run; each run gets its own private browser
context (fresh cookies and storage, closed when the run ends). Chromium is
started by us, as an unprivileged user when the app runs as root, and Playwright
attaches to it over a local DevTools pipe (`connect_over_cdp`), so a page that
breaks out of the renderer still cannot read the app's environment.

Every request goes through a guard: no private or loopback addresses, no
sign-in or payment pages, no form posts that navigate (see web_plan.page_blocked).
Playwright's async API runs on its own event loop thread; the sync helpers below
are what the app calls.
"""

import asyncio
import base64
import html
import ipaddress
import json
import os
import pwd
import re
import shutil
import subprocess
import sys
import tempfile
import threading
import time
import urllib.parse
import urllib.request
from pathlib import Path
from urllib.parse import urlsplit

from playwright.async_api import async_playwright

from web_plan import key_name, page_blocked

VIEWPORT = {"width": 1280, "height": 800}
# Bing: Google search and DuckDuckGo answer an automated browser with a robot check, and we do not
# disguise the browser to get past those, so the agent starts where it is welcome.
START_URL = os.environ.get("WEB_START_URL", "https://www.bing.com/?setlang=en&cc=us")
MAX_PAGES = 4  # open tabs per run; the oldest extra tab is closed

# Hugging Face's own public hosts. Inside a Space they resolve to private addresses, which the
# private-address guard would otherwise block.
_HF_PUBLIC = re.compile(r"^(?:[\w-]+\.)*(?:huggingface\.co|hf\.co)$", re.I)

# Public demo: Bing's SafeSearch is always Strict.
SAFE_SEARCH_COOKIES = [{"name": "SRCHHPGUSR", "value": "ADLT=STRICT", "domain": ".bing.com", "path": "/"}]
# Cloudflare's family resolver answers 0.0.0.0 for adult and malware domains; pages on them are not opened.
FAMILY_DNS = "https://family.cloudflare-dns.com/dns-query"

# Links and pop-ups open in the same tab, so the agent keeps one page and `back` works.
SAME_TAB_JS = """
document.addEventListener('click', e => {
  const a = e.target && e.target.closest && e.target.closest('a[target]');
  if (a && !['_self', '_top', '_parent'].includes(a.target)) a.target = '_self';
}, true);
"""

BLOCK_PAGE = """<!doctype html><meta charset="utf-8"><title>Not in this demo</title>
<body style="margin:0;font:16px system-ui,sans-serif;background:#FBF6EE;color:#2B2722;display:flex;align-items:center;
justify-content:center;height:100vh"><div style="max-width:560px;padding:32px;background:#fff;border:1px solid #E6E3DF;
border-radius:14px;box-shadow:0 6px 22px rgba(60,40,20,.10)"><div style="font-size:13px;letter-spacing:.08em;
text-transform:uppercase;color:#C0552F;font-weight:700">DeskForge demo</div><h1 style="font-size:24px;margin:8px 0 10px">
This page is not available in the demo</h1><p style="line-height:1.55;margin:0 0 10px">Stopped because {reason}.</p>
<p style="line-height:1.55;margin:0;color:#6f6a62;font-size:14px;word-break:break-all">{url}</p></div></body>"""


def ensure_chromium():
    """Install Playwright's Chromium once (into PLAYWRIGHT_BROWSERS_PATH), outside the Xfce desktop's menus."""
    os.environ.setdefault("PLAYWRIGHT_BROWSERS_PATH", "/opt/ms-playwright")
    from playwright.sync_api import sync_playwright

    with sync_playwright() as p:
        if Path(p.chromium.executable_path).exists():
            return
    subprocess.run([sys.executable, "-m", "playwright", "install", "chromium"], check=True)
    os.chmod(os.environ["PLAYWRIGHT_BROWSERS_PATH"], 0o755)


def _browser_user():
    """The unprivileged account Chromium runs as, when we are root (None otherwise)."""
    if os.geteuid() != 0:
        return None
    name = os.environ.get("WEB_BROWSER_USER", "nobody")
    try:
        return pwd.getpwnam(name)
    except KeyError:
        return None


class Session:
    """One run's private browser context and its active tab."""

    def __init__(self, host, context):
        self.host, self.context = host, context
        self.pages = []
        self.blocked = ""  # why the last main-frame navigation was blocked, if it was
        self._adopting = []
        context.on("page", self._on_page)

    # -- tabs -----------------------------------------------------------------

    @property
    def page(self):
        live = [p for p in self.pages if not p.is_closed()]
        self.pages = live
        return live[-1] if live else None

    def _on_page(self, page):
        if page in self.pages:
            return
        self.pages.append(page)
        page.on("dialog", lambda d: asyncio.ensure_future(d.dismiss()))
        page.on("download", lambda d: asyncio.ensure_future(d.cancel()))
        page.on("filechooser", lambda f: None)  # intercepted, never answered
        extra = [p for p in self.pages if not p.is_closed()][:-MAX_PAGES]
        for p in extra:
            asyncio.ensure_future(p.close())
        self._adopting.append(asyncio.ensure_future(self._adopt(page)))

    async def _adopt(self, popup):
        """A pop-up (window.open) becomes a navigation of the tab that opened it."""
        try:
            opener = await popup.opener()
        except Exception:
            return
        if opener is None or opener.is_closed():
            return
        try:
            await popup.wait_for_load_state("commit", timeout=6000)
        except Exception:
            pass
        url = popup.url
        await popup.close()
        if url and url != "about:blank":
            try:
                await opener.goto(url, wait_until="domcontentloaded", timeout=15000)
            except Exception:
                pass

    # -- guard ------------------------------------------------------------------

    async def guard(self, route, request):
        parts = urlsplit(request.url)
        if parts.scheme not in ("http", "https"):
            return await route.continue_()
        host = (parts.hostname or "").lower()
        if await self.host.is_private(host):
            return await route.abort("blockedbyclient")
        if request.is_navigation_request():
            reason = page_blocked(host, parts.path, request.method)
            if not reason and await self.host.is_unsafe(host):
                reason = "this site is blocked in the demo (adult or unsafe content)"
            if reason:
                if request.frame.parent_frame is None:
                    self.blocked = reason
                    body = BLOCK_PAGE.format(reason=html.escape(reason), url=html.escape(request.url[:300]))
                    return await route.fulfill(status=200, content_type="text/html", body=body)
                return await route.abort("blockedbyclient")
        await route.continue_()

    async def guard_ws(self, ws):
        host = (urlsplit(ws.url).hostname or "").lower()
        if await self.host.is_private(host):
            await ws.close()
        else:
            ws.connect_to_server()

    # -- actions ----------------------------------------------------------------

    async def open(self, url):
        page = await self.context.new_page()
        if page not in self.pages:
            self._on_page(page)
        await page.goto(url, wait_until="domcontentloaded", timeout=20000)
        await self.settle(page, first=True)

    async def settle(self, page=None, first=False):
        """Let the page finish what the last action started (navigation, rendering)."""
        await asyncio.sleep(0.35)
        while self._adopting:
            try:
                await asyncio.wait_for(self._adopting.pop(), timeout=10)
            except Exception:
                pass
        page = self.page or page
        if page is None:
            return
        try:
            await page.wait_for_load_state("domcontentloaded", timeout=8000)
        except Exception:
            pass
        try:
            await page.wait_for_load_state("networkidle", timeout=1500 if first else 800)
        except Exception:
            pass
        await asyncio.sleep(0.25)

    async def screenshot(self):
        page = self.page
        if page is None:
            page = await self.context.new_page()
            await page.goto("about:blank")
        try:
            await asyncio.wait_for(page.bring_to_front(), 3)  # a background tab may never paint the next frame
            return await page.screenshot(type="png", timeout=8000, caret="initial")
        except Exception:
            cdp = await self.context.new_cdp_session(page)  # plain capture, without waiting for fonts or stable frames
            try:
                shot = await asyncio.wait_for(cdp.send("Page.captureScreenshot", {"format": "png", "fromSurface": True}), 8)
                return base64.b64decode(shot["data"])
            finally:
                await cdp.detach()

    async def info(self):
        page = self.page
        if page is None:
            return {"url": "", "title": ""}
        try:
            title = await asyncio.wait_for(page.title(), 3)  # runs in the page: a busy page may not answer
        except Exception:
            title = ""
        return {"url": page.url, "title": title}

    async def recover(self, url=START_URL):
        """Replace a tab that stopped responding (e.g. a heavy robot-check page) with a fresh one."""
        for page in list(self.pages):
            try:
                await asyncio.wait_for(page.close(), 5)
            except Exception:
                pass
        self.pages = []
        await self.open(url)

    async def act(self, plan, points):
        """Run one planner action. `points` are pixel positions for target/destination. Returns an error or ''."""
        page = self.page
        self.blocked = ""
        if page is None:
            return "no open page"
        a = plan["action"]
        m, kb = page.mouse, page.keyboard
        if a in ("click", "double_click", "right_click"):
            x, y = points[0]
            await m.move(x, y, steps=4)
            await m.click(x, y, button="right" if a == "right_click" else "left",
                          click_count=2 if a == "double_click" else 1, delay=40)
        elif a == "hover":
            await m.move(*points[0], steps=6)
        elif a == "drag":
            await m.move(*points[0])
            await m.down()
            try:
                await m.move(*points[1], steps=12)
            finally:
                await m.up()
        elif a == "type":
            if plan["clear"]:
                await kb.press("Control+a")
            await kb.type(plan["text"], delay=12)
            if plan.get("submit"):
                await kb.press("Enter")
        elif a == "press":
            await kb.press("+".join(key_name(k) for k in plan["keys"]))
        elif a == "scroll":
            await m.move(*points[0])
            n = plan["amount"] * 120
            dx, dy = {"up": (0, -n), "down": (0, n), "left": (-n, 0), "right": (n, 0)}[plan["direction"]]
            await m.wheel(dx, dy)
            await asyncio.sleep(0.35)
        elif a == "goto":
            try:
                await page.goto(plan["url"].strip(), wait_until="domcontentloaded", timeout=20000)
            except Exception as exc:
                return f"could not open the address ({type(exc).__name__})"
        elif a == "back":
            before = page.url
            response = await page.go_back(wait_until="domcontentloaded", timeout=10000)
            if response is None and page.url == before:  # in-page (SPA) history returns no response but changes the URL
                return "there is no previous page"
        elif a == "wait":
            await asyncio.sleep(1.5)
        await self.settle()
        return ""

    async def close(self):
        try:
            await self.context.close()
        except Exception:
            pass


class BrowserHost:
    """Owns the Chromium process and the event loop thread Playwright runs on."""

    def __init__(self, display=None):
        self.display = display  # an X display for a headed browser; None = headless
        self.loop = asyncio.new_event_loop()
        threading.Thread(target=self.loop.run_forever, daemon=True, name="web-browser").start()
        self._pw = self._browser = self._proc = None
        self._lock = None
        self._dns = {}
        self._family = {}
        self.root = Path(tempfile.mkdtemp(prefix="dfw-"))  # short: Chromium's socket path must stay under 108 bytes
        self._profiles = 0
        os.chmod(self.root, 0o755)

    def call(self, coro, timeout=90):
        return asyncio.run_coroutine_threadsafe(coro, self.loop).result(timeout)

    # -- adult/malware check (family DNS) ------------------------------------------

    async def _family_answer(self, host):
        def ask():
            req = urllib.request.Request(f"{FAMILY_DNS}?name={urllib.parse.quote(host)}&type=A",
                                         headers={"accept": "application/dns-json"})
            with urllib.request.urlopen(req, timeout=3) as r:
                return json.load(r)

        return await asyncio.wait_for(self.loop.run_in_executor(None, ask), 4)

    async def is_unsafe(self, host):
        """True when Cloudflare's family resolver blocks the host. Fails open: no answer means allowed."""
        if not host:
            return False
        hit = self._family.get(host)
        if hit and time.time() - hit[1] < 600:
            return hit[0]
        try:
            answer = await self._family_answer(host)
            bad = any(a.get("data") in ("0.0.0.0", "::") for a in answer.get("Answer", []))
        except Exception:
            bad = False
        self._family[host] = (bad, time.time())
        return bad

    async def family_dns_ok(self):
        """Whether the family resolver answers at all (logged at start-up; when it does not, pages are not filtered)."""
        try:
            return bool((await self._family_answer("example.com")).get("Answer"))
        except Exception:
            return False

    # -- private-address check (SSRF guard) -------------------------------------

    async def is_private(self, host):
        if not host:
            return True
        if _HF_PUBLIC.search(host):  # inside a Space these resolve to private addresses, but they are the public site
            return False
        hit = self._dns.get(host)
        if hit and time.time() - hit[1] < 300:
            return hit[0]
        if host == "localhost" or host.endswith((".localhost", ".local", ".internal", ".lan")):
            bad = True
        else:
            try:
                ips = [ipaddress.ip_address(host.strip("[]"))]
            except ValueError:
                try:
                    infos = await self.loop.getaddrinfo(host, None)
                    ips = [ipaddress.ip_address(i[4][0].split("%")[0]) for i in infos]
                except OSError:
                    ips = []
            bad = any(not ip.is_global or ip.is_multicast for ip in ips)
        self._dns[host] = (bad, time.time())
        return bad

    # -- browser process ----------------------------------------------------------

    async def _ensure(self):
        if self._lock is None:
            self._lock = asyncio.Lock()
        async with self._lock:
            if self._browser is not None and self._browser.is_connected():
                return self._browser
            await self._shutdown()
            self._pw = await async_playwright().start()
            exe = os.environ.get("WEB_CHROMIUM") or self._pw.chromium.executable_path
            user = _browser_user()
            self._profiles += 1
            profile = self.root / f"p{self._profiles}"
            profile.mkdir()
            if user:
                os.chown(profile, user.pw_uid, user.pw_gid)
            args = [exe, "--remote-debugging-port=0", f"--user-data-dir={profile}", "--no-first-run",
                    "--no-default-browser-check", "--disable-dev-shm-usage", "--disable-sync",
                    "--disable-extensions", "--disable-component-update", "--password-store=basic",
                    "--use-mock-keychain", "--lang=en-US", "--window-position=0,0",
                    f"--window-size={VIEWPORT['width']},{VIEWPORT['height'] + 120}", "about:blank"]
            args.insert(1, "--no-sandbox")  # as Playwright does by default: no user namespaces here or in the Space
            env = {"HOME": str(profile), "PATH": "/usr/local/bin:/usr/bin:/bin", "LANG": "en_US.UTF-8",
                   "TMPDIR": str(profile)}
            if self.display:
                env["DISPLAY"] = self.display
            else:
                args.insert(1, "--headless=new")
            log = open(self.root / "chromium.log", "ab")
            self._proc = subprocess.Popen(args, env=env, stdout=log, stderr=log,
                                          user=user.pw_uid if user else None, group=user.pw_gid if user else None,
                                          start_new_session=True)
            port_file = profile / "DevToolsActivePort"
            for _ in range(200):
                if port_file.exists() and port_file.read_text().strip():
                    break
                if self._proc.poll() is not None:
                    tail = (self.root / "chromium.log").read_bytes()[-600:].decode(errors="replace")
                    raise RuntimeError(f"Chromium did not start: {tail}")
                await asyncio.sleep(0.1)
            port = int(port_file.read_text().split()[0])
            self._browser = await self._pw.chromium.connect_over_cdp(f"http://127.0.0.1:{port}")
            return self._browser

    async def _shutdown(self):
        for closer in (getattr(self._browser, "close", None), getattr(self._pw, "stop", None)):
            if closer:
                try:
                    await closer()
                except Exception:
                    pass
        if self._proc and self._proc.poll() is None:
            self._proc.kill()
        self._pw = self._browser = self._proc = None

    async def new_session(self, url=START_URL):
        browser = await self._ensure()
        context = await browser.new_context(viewport=VIEWPORT, device_scale_factor=1, locale="en-US",
                                            accept_downloads=False, service_workers="block",
                                            color_scheme="light")
        session = Session(self, context)
        await context.add_cookies(SAFE_SEARCH_COOKIES)
        await context.add_init_script(SAME_TAB_JS)
        await context.route("**/*", session.guard)
        if hasattr(context, "route_web_socket"):
            await context.route_web_socket("**/*", session.guard_ws)
        try:
            await session.open(url)
        except Exception:
            pass  # a slow start page still leaves a usable tab; the planner sees what loaded
        return session

    def cleanup(self):
        try:
            self.call(self._shutdown(), timeout=20)
        finally:
            shutil.rmtree(self.root, ignore_errors=True)