#!/usr/bin/env python3
"""Website screenshot capture tool — runs INSIDE the playwright/python container.

Usage (inside container): python /work/capture.py /work/sites.json /work/out

Reads a JSON array of {"slug": "kebab-case", "url": "https://..."} and writes
two full-page PNGs per site into the output dir:
  <slug>-desktop.png  (1440x900, device_scale_factor 1)
  <slug>-mobile.png   (390x844, device_scale_factor 2, mobile emulation)

Pure stdlib + playwright (sync API). No other dependencies.
"""

import json
import re
import sys
import time
from pathlib import Path

from playwright.sync_api import sync_playwright, TimeoutError as PWTimeoutError

SLUG_RE = re.compile(r"^[a-z0-9][a-z0-9-]*$")

DESKTOP_UA = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
    "(KHTML, like Gecko) Chrome/128.0.0.0 Safari/537.36"
)
MOBILE_UA = (
    "Mozilla/5.0 (iPhone; CPU iPhone OS 17_5 like Mac OS X) "
    "AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.5 Mobile/15E148 "
    "Safari/604.1"
)

CONSENT_SELECTORS = [
    "#onetrust-accept-btn-handler",
    "button#hs-eu-confirmation-button",
    '[aria-label*="accept" i]',
    'button:has-text("Accept all")',
    'button:has-text("Accept")',
    'button:has-text("I agree")',
    'button:has-text("Got it")',
    ".cc-allow",
    '[data-testid="cookie-accept"]',
]

NAV_TIMEOUT_MS = 30_000
IDLE_TIMEOUT_MS = 15_000
SETTLE_SECONDS = 2.5
SCROLL_STEPS = 8
MAX_ATTEMPTS = 2


def load_sites(path: Path):
    try:
        raw = json.loads(path.read_text())
    except Exception as exc:
        print(f"ERROR: could not read/parse {path}: {exc}")
        sys.exit(2)

    if not isinstance(raw, list):
        print("ERROR: sites.json must be a JSON array")
        sys.exit(2)

    sites = []
    for i, entry in enumerate(raw):
        if not isinstance(entry, dict):
            print(f"ERROR: entry {i} is not an object, skipping")
            continue
        slug = entry.get("slug")
        url = entry.get("url")
        if not isinstance(slug, str) or not SLUG_RE.match(slug):
            print(f"ERROR: entry {i} has invalid slug {slug!r}, skipping")
            continue
        if not isinstance(url, str) or not re.match(r"^https?://", url):
            print(f"ERROR: entry {i} ({slug}) has invalid url {url!r}, skipping")
            continue
        sites.append({"slug": slug, "url": url})
    return sites


def dismiss_consent(page):
    for sel in CONSENT_SELECTORS:
        try:
            locator = page.locator(sel).first
            locator.click(timeout=1200)
            time.sleep(0.6)
            return True
        except Exception:
            continue
    return False


def lazy_load_pass(page):
    try:
        viewport_height = page.evaluate("window.innerHeight") or 900
    except Exception:
        viewport_height = 900

    for _ in range(SCROLL_STEPS):
        try:
            page.evaluate(f"window.scrollBy(0, {viewport_height})")
        except Exception:
            pass
        time.sleep(0.25)

    time.sleep(1.0)

    try:
        page.evaluate("window.scrollTo(0, 0)")
    except Exception:
        pass
    time.sleep(0.5)


def capture_one(browser, slug, url, form_factor, out_dir):
    """Attempt a single capture (one form factor) with up to MAX_ATTEMPTS tries."""
    last_reason = "unknown error"

    for attempt in range(1, MAX_ATTEMPTS + 1):
        context = None
        try:
            if form_factor == "desktop":
                context = browser.new_context(
                    viewport={"width": 1440, "height": 900},
                    device_scale_factor=1,
                    user_agent=DESKTOP_UA,
                    locale="en-CA",
                    timezone_id="America/Vancouver",
                )
            else:
                context = browser.new_context(
                    viewport={"width": 390, "height": 844},
                    device_scale_factor=2,
                    is_mobile=True,
                    has_touch=True,
                    user_agent=MOBILE_UA,
                    locale="en-CA",
                    timezone_id="America/Vancouver",
                )

            page = context.new_page()
            page.goto(url, wait_until="domcontentloaded", timeout=NAV_TIMEOUT_MS)

            try:
                page.wait_for_load_state("networkidle", timeout=IDLE_TIMEOUT_MS)
            except PWTimeoutError:
                pass
            except Exception:
                pass

            time.sleep(SETTLE_SECONDS)

            try:
                dismiss_consent(page)
            except Exception:
                pass

            try:
                lazy_load_pass(page)
            except Exception:
                pass

            out_path = out_dir / f"{slug}-{form_factor}.png"
            page.screenshot(path=str(out_path), full_page=True)

            context.close()
            return True, None

        except Exception as exc:
            last_reason = f"{type(exc).__name__}: {exc}"
            try:
                if context is not None:
                    context.close()
            except Exception:
                pass
            if attempt < MAX_ATTEMPTS:
                time.sleep(1.0)
                continue

    return False, last_reason


def main():
    if len(sys.argv) != 3:
        print("usage: capture.py <sites.json> <out-dir>")
        sys.exit(2)

    sites_path = Path(sys.argv[1])
    out_dir = Path(sys.argv[2])
    out_dir.mkdir(parents=True, exist_ok=True)

    sites = load_sites(sites_path)
    if not sites:
        print("ERROR: no valid sites to capture")
        sys.exit(1)

    results = []  # (slug, form_factor, ok, reason)

    with sync_playwright() as p:
        browser = p.chromium.launch(headless=True)
        try:
            for site in sites:
                slug = site["slug"]
                url = site["url"]
                for form_factor in ("desktop", "mobile"):
                    ok, reason = capture_one(browser, slug, url, form_factor, out_dir)
                    if ok:
                        print(f"{slug} {form_factor}: ok")
                    else:
                        print(f"{slug} {form_factor}: failed ({reason})")
                    results.append((slug, form_factor, ok, reason))
        finally:
            browser.close()

    print()
    print("Summary")
    print("-------")
    width_slug = max((len(r[0]) for r in results), default=4)
    for slug, form_factor, ok, reason in results:
        status = "OK" if ok else "FAILED"
        line = f"{slug:<{width_slug}}  {form_factor:<8}  {status}"
        if not ok:
            line += f"  ({reason})"
        print(line)

    total = len(results)
    passed = sum(1 for r in results if r[2])
    print()
    print(f"{passed}/{total} captures succeeded")

    if passed == total:
        sys.exit(0)
    else:
        sys.exit(1)


if __name__ == "__main__":
    main()