#!/usr/bin/env python3
"""Nadir crawler: render sites light + emulated dark, extract features, classify dark-mode support.

Output per page: OUT/pages/<id>/{light.json.gz, dark.json.gz?, shot.jpg, meta.json}
and a line in OUT/index.jsonl.
"""
import argparse, asyncio, base64, gzip, hashlib, json, os, random, re, sys, time, traceback
from pathlib import Path
from playwright.async_api import async_playwright

HERE = Path(__file__).resolve().parent
EXTRACT_JS = (HERE / "extract.js").read_text()

NO_ANIM = """(() => { try { const s = new CSSStyleSheet(); s.replaceSync('*,*::before,*::after{transition:none!important;animation-duration:0s!important;animation-delay:0s!important;caret-color:auto}');
 document.adoptedStyleSheets = [...document.adoptedStyleSheets, s]; window.__nadirNoAnim = s; } catch (e) {} })()"""
DROP_NO_ANIM = """(() => { try { document.adoptedStyleSheets = document.adoptedStyleSheets.filter(x => x !== window.__nadirNoAnim); } catch (e) {} })()"""

FETCH_SHIM = """window.__nadirFetchImage = async (src) => {
  try { const r = await window.__nadirFetchB64(src); if (!r) return null;
    const bin = atob(r.b64); const u = new Uint8Array(bin.length); for (let i = 0; i < bin.length; i++) u[i] = bin.charCodeAt(i);
    return new Blob([u], {type: r.ct || 'application/octet-stream'}); } catch (e) { return null; } };"""


def page_id(url):
    return hashlib.sha1(url.encode()).hexdigest()[:16]


async def settle(page, ms=600):
    try:
        await page.wait_for_timeout(ms)
    except Exception:
        pass


async def autoscroll(page):
    # trigger lazy content, then return to top
    try:
        await page.evaluate("""async () => {
          const H = Math.min(document.documentElement.scrollHeight, innerHeight * 6);
          for (let y = 0; y < H; y += innerHeight * 0.8) { scrollTo(0, y); await new Promise(r => setTimeout(r, 120)); }
          scrollTo(0, 0); await new Promise(r => setTimeout(r, 250)); }""")
    except Exception:
        pass


async def extract(page, **opts):
    await page.evaluate(EXTRACT_JS)
    o = dict(opts)
    js = """async (o) => await __nadirExtract(o)"""
    return await page.evaluate(js, o)


DARK_TOGGLE_JS = """(sel) => {
  const html = document.documentElement, body = document.body;
  const undo = [];
  const m = sel.match(/^(html|body|:root)?(.*)$/);
  const target = (m[1] === 'body') ? body : html;
  const rest = m[2];
  let cm = rest.match(/^\\.([\\w-]+)$/);
  let am = rest.match(/^\\[([\\w-]+)(?:[~|^*$]?=\\s*["']?([\\w-]+)["']?)?\\s*(?:i\\s*)?\\]$/);
  if (cm) { for (const t of [html, body]) { if (!t.classList.contains(cm[1])) { t.classList.add(cm[1]); undo.push(['c', t, cm[1]]); } } }
  else if (am) { const v = am[2] || 'dark'; for (const t of [html, body]) { const old = t.getAttribute(am[1]); t.setAttribute(am[1], v); undo.push(['a', t, am[1], old]); } }
  else return false;
  window.__nadirUndo = undo; return true; }"""
UNDO_JS = """() => { for (const u of (window.__nadirUndo || [])) { if (u[0] === 'c') u[1].classList.remove(u[2]); else if (u[3] === null) u[1].removeAttribute(u[2]); else u[1].setAttribute(u[2], u[3]); } }"""


LINKS_JS = r"""() => { const o = location.origin, seen = new Set(), out = [];
  for (const a of document.querySelectorAll('a[href]')) { let u; try { u = new URL(a.href); } catch (e) { continue; }
    if (u.origin !== o || u.pathname.length < 2 || /\.(pdf|zip|png|jpe?g|gif|svg|mp4|xml|json)$/i.test(u.pathname) || /login|signin|signup|register|cart|checkout|account|logout|auth/i.test(u.pathname)) continue;
    u.hash = ''; const k = u.pathname; if (seen.has(k)) continue; seen.add(k); out.push(u.href); if (out.length >= 12) break; }
  return out; }"""


IMG_STATS_JS = """async (a) => {
  const res = await __nadirImageStats(a.images, {shot: a.shot, bigThumbSize: 96, concurrency: 8});
  a.images.forEach((o, i) => { Object.assign(o, res[i]); o.f = __nadirImageFeatures(o, a.vw, a.vh, a.docH); });
  return a.images; }"""


async def image_stats(ctx, d, shot_png):
    """Pixel stats in a separate blank page (like the extension's offscreen document): no page CSP, no Python transfer."""
    if not d["images"]:
        return
    helper = await ctx.new_page()
    try:
        await helper.evaluate(EXTRACT_JS)
        shot = "data:image/png;base64," + base64.b64encode(shot_png).decode()
        d["images"] = await asyncio.wait_for(helper.evaluate(IMG_STATS_JS, {"images": d["images"], "shot": shot, "vw": d["vw"], "vh": d["vh"], "docH": d["docH"]}), 40)
    except Exception as e:
        d["imgStatsError"] = str(e)[:200]
    finally:
        await helper.close()


def is_dark(d):
    dk = d["darkness"]
    return dk["bgL"] < 0.45 and dk["darkFrac"] > 0.6 and (dk["textL"] == 0 or dk["textL"] > 0.55)


async def crawl_one(ctx, url, out_dir, args):
    pid = page_id(url)
    pdir = out_dir / "pages" / pid
    rec = {"id": pid, "url": url, "t": time.time()}
    page = await ctx.new_page()
    try:
        t0 = time.time()
        resp = await page.goto(url, wait_until="domcontentloaded", timeout=args.nav_timeout * 1000)
        rec["status"] = resp.status if resp else None
        rec["final_url"] = page.url
        try:
            await page.wait_for_load_state("load", timeout=8000)
        except Exception:
            pass
        await settle(page, 1200)
        await page.evaluate(NO_ANIM)
        await autoscroll(page)
        await settle(page, 700)
        light = await extract(page, describe=True, thumbs=False, elemColors=True, markSvg=True, maxImages=args.max_images)
        shot_png = await page.screenshot(type="png")
        await image_stats(ctx, light, shot_png)
        rec["light_ms"] = light["stats"]["ms"]
        rec["n_tokens"] = len(light["tokens"])
        rec["n_images"] = len(light["images"])
        rec["darkness"] = light["darkness"]
        rec["scheme"] = light["scheme"]
        if light["stats"]["visible"] < 15 or len(light["tokens"]) < 4:
            rec["kind"] = "empty"
            return rec
        pdir.mkdir(parents=True, exist_ok=True)
        # screenshot (top ~2 screens)
        H = min(int(light["docH"]), args.shot_h)
        await page.screenshot(path=str(pdir / "shot.jpg"), type="jpeg", quality=80, full_page=True,
                              clip={"x": 0, "y": 0, "width": 1280, "height": H})
        light_dark = is_dark(light)
        kind = "dark" if light_dark else "light"
        dark = None
        # 1) emulated prefers-color-scheme: dark
        await page.emulate_media(color_scheme="dark")
        await settle(page, 1500)
        dm = await extract(page, describe=False, thumbs=False, elemColors=True, pseudo=True)
        if not light_dark and is_dark(dm):
            kind, dark = "native_media", dm
        elif light_dark:
            pass
        else:
            # 2) class / attribute toggles found in stylesheets
            await page.emulate_media(color_scheme="light")
            await settle(page, 400)
            for sel, n in (light["scheme"]["darkSelectors"] or [])[:3]:
                if n < 8:
                    continue
                ok = await page.evaluate(DARK_TOGGLE_JS, sel)
                if not ok:
                    continue
                await settle(page, 900)
                dc = await extract(page, describe=False, thumbs=False, elemColors=True)
                if is_dark(dc):
                    kind, dark = "native_class", dc
                    rec["toggle"] = sel
                    break
                await page.evaluate(UNDO_JS)
                await settle(page, 300)
        rec["kind"] = kind
        try:
            rec["links"] = await page.evaluate(LINKS_JS)
        except Exception:
            rec["links"] = []
        if dark is not None:
            rec["dark_darkness"] = dark["darkness"]
            await page.screenshot(path=str(pdir / "shot_dark.jpg"), type="jpeg", quality=80, full_page=True,
                                  clip={"x": 0, "y": 0, "width": 1280, "height": H})
            with gzip.open(pdir / "dark.json.gz", "wt") as f:
                json.dump(dark, f, separators=(",", ":"))
        with gzip.open(pdir / "light.json.gz", "wt") as f:
            json.dump(light, f, separators=(",", ":"))
        rec["secs"] = round(time.time() - t0, 1)
        return rec
    except Exception as e:
        rec["kind"] = "error"
        rec["error"] = (type(e).__name__ + ": " + str(e))[:300]
        return rec
    finally:
        try:
            await page.close()
        except Exception:
            pass


async def worker(wid, browser_pool, queue, out_dir, args, index_f, counters):
    while True:
        try:
            url = queue.get_nowait()
        except asyncio.QueueEmpty:
            return
        browser = browser_pool[wid % len(browser_pool)]
        ctx = None
        try:
            ctx = await browser.new_context(viewport={"width": 1280, "height": 800}, device_scale_factor=1,
                                            color_scheme="light", locale="en-US", ignore_https_errors=True,
                                            user_agent=args.ua, java_script_enabled=True, bypass_csp=True)
            ctx.set_default_timeout(15000)
            rec = await asyncio.wait_for(crawl_one(ctx, url, out_dir, args), timeout=args.page_timeout)
        except asyncio.TimeoutError:
            rec = {"id": page_id(url), "url": url, "kind": "error", "error": "page timeout"}
        except Exception as e:
            rec = {"id": page_id(url), "url": url, "kind": "error", "error": str(e)[:300]}
        finally:
            if ctx:
                try:
                    await asyncio.wait_for(ctx.close(), 10)
                except Exception:
                    pass
        counters[rec.get("kind", "error")] = counters.get(rec.get("kind", "error"), 0) + 1
        index_f.write(json.dumps(rec) + "\n")
        index_f.flush()


async def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--urls", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--workers", type=int, default=24)
    ap.add_argument("--browsers", type=int, default=4)
    ap.add_argument("--nav-timeout", type=int, default=25)
    ap.add_argument("--page-timeout", type=int, default=90)
    ap.add_argument("--shot-h", type=int, default=2000)
    ap.add_argument("--max-images", type=int, default=48, help="images kept per page (the model sees the first 48; round 2 keeps more for the element head)")
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--shard", type=int, default=0)
    ap.add_argument("--nshards", type=int, default=1)
    ap.add_argument("--ua", default="Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36")
    args = ap.parse_args()
    out_dir = Path(args.out)
    (out_dir / "pages").mkdir(parents=True, exist_ok=True)
    done = set()
    for ip in out_dir.glob("index*.jsonl"):
        for line in ip.open():
            try:
                done.add(json.loads(line)["url"])
            except Exception:
                pass
    idx = out_dir / ("index.jsonl" if args.nshards == 1 else f"index.{args.shard}.jsonl")
    urls = [u.strip() for u in open(args.urls) if u.strip() and not u.startswith("#")]
    urls = [u for u in urls if u not in done and (args.nshards == 1 or int(hashlib.sha1(u.encode()).hexdigest(), 16) % args.nshards == args.shard)]
    if args.limit:
        urls = urls[: args.limit]
    print(f"{len(urls)} urls to crawl ({len(done)} done)", flush=True)
    q = asyncio.Queue()
    for u in urls:
        q.put_nowait(u)
    counters = {}
    async with async_playwright() as p:
        browsers = []
        for _ in range(args.browsers):
            browsers.append(await p.chromium.launch(headless=True, args=[
                "--disable-web-security", "--disable-features=IsolateOrigins,site-per-process,BlockInsecurePrivateNetworkRequests",
                "--disable-dev-shm-usage", "--no-first-run", "--mute-audio", "--hide-scrollbars", "--font-render-hinting=none",
                "--autoplay-policy=user-gesture-required", "--disable-notifications"]))
        with idx.open("a") as f:
            start = time.time()
            async def report():
                while True:
                    await asyncio.sleep(60)
                    n = sum(counters.values())
                    print(f"[{time.strftime('%H:%M:%S')}] {n} done, {n / max(1, time.time() - start) * 3600:.0f}/h {counters}", flush=True)
            rep = asyncio.create_task(report())
            await asyncio.gather(*[worker(i, browsers, q, out_dir, args, f, counters) for i in range(args.workers)])
            rep.cancel()
        for b in browsers:
            await b.close()
    print("DONE", counters, flush=True)


if __name__ == "__main__":
    asyncio.run(main())