#!/usr/bin/env python3
"""
Design-forensics capture.

Renders a live website (JS included) and pulls down everything you need to
study or rebuild its DESIGN — not its data:

  - full-page screenshots at desktop + mobile breakpoints
  - the fully-rendered HTML (post-JS DOM)
  - the raw <link> stylesheets
  - a ranked design-token report: fonts, colors, font sizes/weights,
    spacing scale, radii — each ordered by how often it actually appears

Usage:
    python capture.py <url> [name]

Example:
    python capture.py https://cofounder.co cofounder
Output lands in ./captures/<name>/
"""
import sys, os, json, re, urllib.parse, urllib.request
from collections import Counter
from playwright.sync_api import sync_playwright

DESKTOP = {"width": 1440, "height": 900}
MOBILE = {"width": 390, "height": 844}

# JS that walks the rendered DOM and tallies the computed design tokens.
EXTRACT_JS = r"""
() => {
  const tally = {
    fontFamily: {}, fontSize: {}, fontWeight: {},
    color: {}, background: {}, letterSpacing: {},
    lineHeight: {}, radius: {}, spacing: {}
  };
  const bump = (bucket, key) => {
    if (!key || key === 'normal' || key === 'none' || key === 'auto') return;
    bucket[key] = (bucket[key] || 0) + 1;
  };
  const isVisibleColor = (c) =>
    c && c !== 'rgba(0, 0, 0, 0)' && c !== 'transparent';

  for (const el of document.querySelectorAll('*')) {
    const s = getComputedStyle(el);
    const txt = (el.textContent || '').trim().length > 0;
    bump(tally.fontFamily, s.fontFamily);
    if (txt) {
      bump(tally.fontSize, s.fontSize);
      bump(tally.fontWeight, s.fontWeight);
      bump(tally.lineHeight, s.lineHeight);
      bump(tally.letterSpacing, s.letterSpacing);
      if (isVisibleColor(s.color)) bump(tally.color, s.color);
    }
    if (isVisibleColor(s.backgroundColor)) bump(tally.background, s.backgroundColor);
    bump(tally.radius, s.borderRadius);
    for (const side of ['marginTop','marginBottom','paddingTop','paddingBottom','gap']) {
      bump(tally.spacing, s[side]);
    }
  }
  // Also grab the declared @font-face / font links and CSS custom props on :root
  const rootVars = {};
  const rs = getComputedStyle(document.documentElement);
  for (const name of rs) {
    if (name.startsWith('--')) rootVars[name] = rs.getPropertyValue(name).trim();
  }
  const stylesheets = [...document.querySelectorAll('link[rel="stylesheet"]')].map(l => l.href);
  const fonts = [...document.querySelectorAll('link[href*="font"], link[href*="Font"]')].map(l => l.href);
  return { tally, rootVars, stylesheets, fonts,
           title: document.title, url: location.href };
}
"""

def rank(d, top=15):
    return [{"value": k, "count": v} for k, v in
            sorted(d.items(), key=lambda kv: -kv[1])[:top]]

def md_table(title, items):
    rows = "\n".join(f"| `{i['value']}` | {i['count']} |" for i in items)
    return f"### {title}\n\n| value | uses |\n|---|---|\n{rows}\n"

def main():
    if len(sys.argv) < 2:
        print(__doc__); sys.exit(1)
    url = sys.argv[1]
    name = sys.argv[2] if len(sys.argv) > 2 else \
        urllib.parse.urlparse(url).netloc.replace("www.", "").split(".")[0]
    outdir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "captures", name)
    os.makedirs(outdir, exist_ok=True)
    print(f"→ capturing {url}\n→ output: {outdir}")

    with sync_playwright() as p:
        browser = p.chromium.launch()
        ctx = browser.new_context(viewport=DESKTOP, device_scale_factor=2,
                                  user_agent=("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
                                              "AppleWebKit/537.36 (KHTML, like Gecko) "
                                              "Chrome/148.0 Safari/537.36"))
        page = ctx.new_page()
        page.goto(url, wait_until="networkidle", timeout=60000)
        page.wait_for_timeout(1500)  # let lazy/animated content settle

        # rendered HTML
        html = page.content()
        with open(os.path.join(outdir, "rendered.html"), "w") as f:
            f.write(html)

        # desktop screenshot
        page.screenshot(path=os.path.join(outdir, "desktop-full.png"), full_page=True)

        # design tokens
        data = page.evaluate(EXTRACT_JS)

        # mobile screenshot
        page.set_viewport_size(MOBILE)
        page.wait_for_timeout(800)
        page.screenshot(path=os.path.join(outdir, "mobile-full.png"), full_page=True)

        browser.close()

    # rank + write tokens.json
    t = data["tally"]
    tokens = {
        "title": data["title"], "url": data["url"],
        "rootVars": data["rootVars"],
        "stylesheets": data["stylesheets"],
        "fontLinks": data["fonts"],
        "fonts": rank(t["fontFamily"]),
        "colors": rank(t["color"]),
        "backgrounds": rank(t["background"]),
        "fontSizes": rank(t["fontSize"]),
        "fontWeights": rank(t["fontWeight"]),
        "lineHeights": rank(t["lineHeight"]),
        "letterSpacing": rank(t["letterSpacing"]),
        "radii": rank(t["radius"]),
        "spacing": rank(t["spacing"], top=20),
    }
    with open(os.path.join(outdir, "tokens.json"), "w") as f:
        json.dump(tokens, f, indent=2)

    # download linked stylesheets
    cssdir = os.path.join(outdir, "css")
    os.makedirs(cssdir, exist_ok=True)
    saved = 0
    for i, href in enumerate(data["stylesheets"][:20]):
        try:
            req = urllib.request.Request(href, headers={"User-Agent": "Mozilla/5.0"})
            css = urllib.request.urlopen(req, timeout=15).read()
            fn = re.sub(r"[^a-zA-Z0-9_.-]", "_", urllib.parse.urlparse(href).path.split("/")[-1] or f"sheet{i}")
            if not fn.endswith(".css"): fn += ".css"
            with open(os.path.join(cssdir, f"{i:02d}_{fn}"), "wb") as f:
                f.write(css)
            saved += 1
        except Exception as e:
            print(f"  ! css fetch failed {href}: {e}")

    # human-readable report
    report = [f"# Design forensics — {tokens['title']}", f"`{tokens['url']}`\n",
              "Screenshots: `desktop-full.png`, `mobile-full.png` · "
              "Rendered DOM: `rendered.html` · "
              f"Stylesheets saved: {saved} in `css/`\n"]
    if tokens["rootVars"]:
        report.append("### CSS custom properties (:root)\n")
        for k, v in list(tokens["rootVars"].items())[:40]:
            report.append(f"- `{k}`: `{v}`")
        report.append("")
    report.append(md_table("Font families", tokens["fonts"]))
    report.append(md_table("Text colors", tokens["colors"]))
    report.append(md_table("Background colors", tokens["backgrounds"]))
    report.append(md_table("Font sizes", tokens["fontSizes"]))
    report.append(md_table("Font weights", tokens["fontWeights"]))
    report.append(md_table("Spacing scale", tokens["spacing"]))
    report.append(md_table("Border radii", tokens["radii"]))
    with open(os.path.join(outdir, "REPORT.md"), "w") as f:
        f.write("\n".join(report))

    print(f"✓ done — {outdir}")
    print(f"  screenshots, rendered.html, tokens.json, REPORT.md, {saved} stylesheets")

if __name__ == "__main__":
    main()