#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Superium — screenshots
======================

Captures every page with headless Chrome (Edge as a fallback) into
``assets/images/shots/``. Used by the showcase and for eyeballing responsive
behaviour.

    python tools/shots.py                 # all pages, desktop + mobile
    python tools/shots.py --only market   # just pages matching "market"
    python tools/shots.py --width 360     # one extra width
    python tools/shots.py --viewport      # crop to the fold instead

Shots are **whole pages**, header to footer, not just the first screenful —
the showcase lightbox shows the page as a buyer would scroll it. Output is
WebP, like every other image in the package.

Two engines, both driving the Chrome that is already installed:

* Playwright, if the module is importable — it can capture past the fold.
  It uses the local browser (``executable_path``), so there is nothing to
  download.
* Otherwise ``chrome --headless --screenshot``, which can only reach as far
  as the window. That path prints a notice and falls back to fold-height
  shots.

A local web server must be running (the script starts one itself if the port
is free).
"""

from __future__ import annotations

import argparse
import http.server
import io
import json
import shutil
import socket
import socketserver
import subprocess
import sys
import threading
import time
from pathlib import Path

try:  # optional — see the module docstring
    from playwright.sync_api import sync_playwright
except ImportError:  # pragma: no cover - depends on the environment
    sync_playwright = None

try:
    from PIL import Image
except ImportError:  # pragma: no cover
    sys.exit("✗ Pillow is required: pip install pillow")

ROOT = Path(__file__).resolve().parent.parent
SHOTS = ROOT / "assets" / "images" / "shots"
PAGES_JSON = ROOT / "assets" / "data" / "pages.json"

PORT = 8791

CHROME_CANDIDATES = [
    r"C:\Program Files\Google\Chrome\Application\chrome.exe",
    r"C:\Program Files (x86)\Google\Chrome\Application\chrome.exe",
    r"C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe",
    r"C:\Program Files\Microsoft\Edge\Application\msedge.exe",
    "/usr/bin/google-chrome",
    "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome",
]

VIEWS = {
    "desktop": (1440, 1000),
    "mobile": (390, 844),
}

# WebP quality for the captures. 86 keeps Persian text crisp at 1:1 while a
# full-page desktop shot stays around a fifth of the PNG's weight.
QUALITY = 86

# Chrome refuses to rasterise past this; also WebP's own ceiling is 16383px.
MAX_HEIGHT = 16000

# These pages render an empty state with no basket, which is not what the
# showcase should advertise. ?demo=1 seeds a sample basket first.
NEEDS_CART = {"cart.html", "vendors.html", "checkout.html", "store.html", "product.html"}

# Floating chrome is hidden while a page is captured: in a full-page shot the
# support button would sit halfway down the image, where it reads as part of
# the page rather than as something pinned to the viewport.
CAPTURE_CSS = ".sp-fabs { display: none !important; }"

# Everything below the fold is lazy, and a full-page capture never scrolls it
# into view — so promote it, walk the page once, then wait for the decode.
EAGER_IMAGES = """() => {
  document.querySelectorAll('img[loading="lazy"]').forEach(function (img) {
    img.loading = 'eager';
  });
}"""

SETTLE = """async () => {
  const step = window.innerHeight;
  for (let y = 0; y < document.body.scrollHeight; y += step) {
    window.scrollTo(0, y);
    await new Promise((done) => setTimeout(done, 60));
  }
  window.scrollTo(0, 0);
  await Promise.all(
    Array.from(document.images)
      .filter((img) => !img.complete)
      .map((img) => new Promise((done) => { img.onload = img.onerror = done; }))
  );
  if (document.fonts && document.fonts.ready) await document.fonts.ready;
}"""


def find_browser() -> str:
    for path in CHROME_CANDIDATES:
        if Path(path).exists():
            return path
    found = shutil.which("chrome") or shutil.which("chromium") or shutil.which("msedge")
    if found:
        return found
    sys.exit(
        "✗ No Chrome or Edge found.\n"
        "  Install one, or add its path to CHROME_CANDIDATES in tools/shots.py"
    )


def port_open(port: int) -> bool:
    with socket.socket() as s:
        s.settimeout(0.4)
        return s.connect_ex(("127.0.0.1", port)) == 0


def serve(port: int) -> socketserver.TCPServer | None:
    """Serve the package root unless something is already on the port."""
    if port_open(port):
        return None

    class Handler(http.server.SimpleHTTPRequestHandler):
        def __init__(self, *a, **kw):
            super().__init__(*a, directory=str(ROOT), **kw)

        def log_message(self, *a):
            pass

    socketserver.TCPServer.allow_reuse_address = True
    httpd = socketserver.TCPServer(("127.0.0.1", port), Handler)
    threading.Thread(target=httpd.serve_forever, daemon=True).start()
    time.sleep(0.6)
    return httpd


def write_webp(png: bytes, dest: Path) -> bool:
    """PNG bytes in, the package's WebP out. Drops a stale PNG of the same name."""
    dest.parent.mkdir(parents=True, exist_ok=True)
    Image.open(io.BytesIO(png)).convert("RGB").save(
        dest, "WEBP", quality=QUALITY, method=5
    )
    legacy = dest.with_suffix(".png")
    if legacy.exists():
        legacy.unlink()
    return dest.exists() and dest.stat().st_size > 1000


def capture_full(page, url: str, dest: Path, width: int, height: int) -> bool:
    """Whole page, header to footer. Playwright only."""
    page.set_viewport_size({"width": width, "height": height})
    page.goto(url, wait_until="load", timeout=60_000)
    page.add_style_tag(content=CAPTURE_CSS)
    page.evaluate(EAGER_IMAGES)
    page.evaluate(SETTLE)
    page.wait_for_timeout(250)

    full = min(page.evaluate("document.documentElement.scrollHeight"), MAX_HEIGHT)
    # Clip to the viewport width: carousel slides sit outside the document box
    # on both sides, and a bare full-page capture would widen the image to
    # include them.
    png = page.screenshot(
        full_page=True, clip={"x": 0, "y": 0, "width": width, "height": full}
    )
    return write_webp(png, dest)


def capture_fold(browser: str, url: str, dest: Path, width: int, height: int) -> bool:
    """As far as the window reaches. The fallback when Playwright is missing."""
    tmp = dest.with_suffix(".tmp.png")
    cmd = [
        browser,
        "--headless=new",
        "--disable-gpu",
        "--hide-scrollbars",
        "--force-device-scale-factor=1",
        "--virtual-time-budget=25000",
        "--run-all-compositor-stages-before-draw",
        f"--window-size={width},{height}",
        f"--screenshot={tmp}",
        url,
    ]
    try:
        subprocess.run(cmd, capture_output=True, timeout=90)
    except subprocess.TimeoutExpired:
        return False
    if not tmp.exists():
        return False
    ok = write_webp(tmp.read_bytes(), dest)
    tmp.unlink()
    return ok


def load_pages() -> list[dict]:
    if not PAGES_JSON.exists():
        sys.exit("✗ assets/data/pages.json missing — run tools/build_pages.py first.")
    return json.loads(PAGES_JSON.read_text(encoding="utf-8"))


def build_jobs(pages: list[dict], views: dict, port: int) -> list[tuple]:
    """(stem, view, url, dest, width, height) for every shot to take."""
    jobs = []
    for entry in pages:
        stem = entry["file"].replace(".html", "")
        for view, (w, h) in views.items():
            query = "?demo=1" if entry["file"] in NEEDS_CART else ""
            jobs.append(
                (
                    stem,
                    view,
                    f"http://127.0.0.1:{port}/{entry['file']}{query}",
                    SHOTS / f"{stem}-{view}.webp",
                    w,
                    h,
                )
            )
    return jobs


def run_full(browser: str, jobs: list[tuple]) -> list[str]:
    """One browser for the whole run — a fresh page per shot keeps state clean."""
    failed = []
    with sync_playwright() as pw:
        chrome = pw.chromium.launch(executable_path=browser)
        for stem, view, url, dest, w, h in jobs:
            page = chrome.new_page(viewport={"width": w, "height": h})
            try:
                if not capture_full(page, url, dest, w, h):
                    failed.append(f"{stem}-{view}")
            except Exception as exc:  # one bad page must not end the run
                print(f"    ! {stem}-{view}: {exc}")
                failed.append(f"{stem}-{view}")
            finally:
                page.close()
            print(f"  ✓ {stem}-{view}")
        chrome.close()
    return failed


def run_fold(browser: str, jobs: list[tuple]) -> list[str]:
    failed = []
    for stem, view, url, dest, w, h in jobs:
        if not capture_fold(browser, url, dest, w, h):
            failed.append(f"{stem}-{view}")
        print(f"  ✓ {stem}-{view}")
    return failed


def main() -> int:
    ap = argparse.ArgumentParser(description="Capture page screenshots")
    ap.add_argument("--only", help="substring filter on the file name")
    ap.add_argument("--width", type=int, help="capture one extra custom width")
    ap.add_argument("--viewport", action="store_true",
                    help="stop at the fold instead of capturing the whole page")
    ap.add_argument("--port", type=int, default=PORT)
    args = ap.parse_args()

    browser = find_browser()
    full_page = sync_playwright is not None and not args.viewport
    print(f"browser: {browser}")
    print(f"mode:    {'full page' if full_page else 'fold only'}")
    if sync_playwright is None and not args.viewport:
        print("         (pip install playwright for whole-page shots)")

    httpd = serve(args.port)
    print(f"serving: http://127.0.0.1:{args.port}  ({'started' if httpd else 'already up'})")

    views = dict(VIEWS)
    if args.width:
        views[f"w{args.width}"] = (args.width, 900)

    pages = load_pages()
    if args.only:
        pages = [p for p in pages if args.only in p["file"]]

    SHOTS.mkdir(parents=True, exist_ok=True)
    jobs = build_jobs(pages, views, args.port)

    try:
        failed = run_full(browser, jobs) if full_page else run_fold(browser, jobs)
    finally:
        if httpd:
            httpd.shutdown()

    print(f"\n{len(jobs) - len(failed)} screenshot(s) → assets/images/shots/")
    if failed:
        print("✗ failed: " + ", ".join(failed))
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())