#!/usr/bin/env python3
"""Visual Intelligence Capture — grab representative images from URLs and analyze them."""

import argparse, json, os, re, sys, time, urllib.request, urllib.parse, urllib.error

GEMINI_KEY = os.environ.get("GEMINI_API_KEY", "AIzaSyC2dTGU667C7l8mzaCjI5S18txWtf0aHPc")
NOTION_KEY_PATH = "/home/clawd/secrets/notion/api_key"

VISION_PROMPT = (
    "You are a senior creative director analyzing a design reference. Describe: "
    "1) What type of creative work this is (brand identity, packaging, web design, editorial, spatial, etc.) "
    "2) Key visual elements: typography approach, color system, layout style, textures "
    "3) What makes this work distinctive or noteworthy "
    "4) One-line 'taste note' that captures the essence. "
    "Be precise and opinionated, not generic. Max 150 words. "
    "Return JSON with keys: work_type, visual_elements, distinctive, taste_note"
)

DEFAULT_HEADERS = {"User-Agent": "Mozilla/5.0 (compatible; VisualCapture/1.0)"}

def _get(url, headers=None, retries=5):
    h = {**DEFAULT_HEADERS, **(headers or {})}
    for i in range(retries):
        try:
            req = urllib.request.Request(url, headers=h)
            with urllib.request.urlopen(req, timeout=30) as r:
                return json.loads(r.read())
        except urllib.error.HTTPError as e:
            if e.code == 429 and i < retries - 1:
                wait = 15 * (2 ** i)  # 15, 30, 60, 120s
                print(f"  Rate limited, waiting {wait}s...", file=sys.stderr)
                time.sleep(wait)
            else:
                raise

def _req(url, data, headers):
    req = urllib.request.Request(url, data=json.dumps(data).encode(), headers=headers, method="POST")
    with urllib.request.urlopen(req, timeout=60) as r:
        return json.loads(r.read())

def _patch(url, data, headers):
    req = urllib.request.Request(url, data=json.dumps(data).encode(), headers=headers, method="PATCH")
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.loads(r.read())


def microlink_capture(url: str) -> dict:
    api = f"https://api.microlink.io/?url={urllib.parse.quote(url, safe='')}&screenshot=true"
    resp = _get(api)
    d = resp.get("data", {})
    img = d.get("image", {})
    ss = d.get("screenshot", {})
    return {
        "title": d.get("title", ""),
        "image_url": img.get("url", "") if isinstance(img, dict) else "",
        "screenshot_url": ss.get("url", "") if isinstance(ss, dict) else "",
        "description": d.get("description", ""),
    }


def assess_image_quality(image_url: str) -> str:
    if not image_url:
        return "garbage"
    low = image_url.lower()
    if low.endswith(".svg") or any(k in low for k in ["logo", "icon", "favicon"]):
        return "logo"
    if any(k in low for k in ["placeholder", "default-og", "1x1.gif", "spacer"]):
        return "garbage"
    return "good"


def _can_download(url: str) -> bool:
    """Check if image URL is actually downloadable."""
    if not url:
        return False
    try:
        req = urllib.request.Request(url, method="HEAD", headers=DEFAULT_HEADERS)
        with urllib.request.urlopen(req, timeout=10) as r:
            return r.status == 200
    except:
        return False


def best_image(cap: dict) -> str:
    quality = assess_image_quality(cap["image_url"])
    if quality == "good":
        # Prefer original but fall back to screenshot if not downloadable
        if _can_download(cap["image_url"]):
            return cap["image_url"]
    return cap["screenshot_url"] or cap["image_url"]


def analyze_with_vision(image_url: str, gemini_key: str = None) -> dict:
    key = gemini_key or GEMINI_KEY
    api = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent?key={key}"
    body = {
        "contents": [{"parts": [
            {"text": VISION_PROMPT},
            {"inline_data": None},  # placeholder replaced below
        ]}]
    }
    # Use image URL via file_data
    body["contents"][0]["parts"][1] = {"file_data": {"mime_type": "image/jpeg", "file_uri": image_url}}
    # Gemini doesn't support arbitrary URLs via file_data; download and inline
    try:
        req = urllib.request.Request(image_url, headers=DEFAULT_HEADERS)
        with urllib.request.urlopen(req, timeout=30) as r:
            img_bytes = r.read()
            ct = r.headers.get("Content-Type", "image/jpeg")
        import base64
        b64 = base64.b64encode(img_bytes).decode()
        body["contents"][0]["parts"][1] = {"inline_data": {"mime_type": ct.split(";")[0], "data": b64}}
    except Exception as e:
        return {"error": str(e)}

    try:
        resp = _req(api, body, {"Content-Type": "application/json"})
        text = resp["candidates"][0]["content"]["parts"][0]["text"]
        # Try to parse JSON from response
        m = re.search(r'\{[^{}]*\}', text, re.DOTALL)
        if m:
            return json.loads(m.group())
        return {"taste_note": text}
    except Exception as e:
        return {"error": str(e)}


def get_notion_key():
    try:
        return open(NOTION_KEY_PATH).read().strip()
    except:
        return os.environ.get("NOTION_API_KEY", "")


def set_notion_cover(page_id: str, image_url: str, note: str = "", notion_key: str = None) -> bool:
    key = notion_key or get_notion_key()
    headers = {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
        "Notion-Version": "2022-06-28",
    }
    base = f"https://api.notion.com/v1/pages/{page_id}"
    data = {"cover": {"type": "external", "external": {"url": image_url}}}
    if note:
        data["properties"] = {
            "Notes": {"rich_text": [{"type": "text", "text": {"content": note[:2000]}}]}
        }
    try:
        _patch(base, data, headers)
        return True
    except urllib.error.HTTPError as e:
        # Try "Note" instead of "Notes"
        if note and "Notes" in data.get("properties", {}):
            data["properties"] = {
                "Note": {"rich_text": [{"type": "text", "text": {"content": note[:2000]}}]}
            }
            try:
                _patch(base, data, headers)
                return True
            except:
                pass
        print(f"  Notion error: {e.read().decode()[:300]}", file=sys.stderr)
        return False


def process_url(url: str, notion_page_id: str = None, analyze: bool = False, output_dir: str = None) -> dict:
    print(f"Capturing: {url}")
    cap = microlink_capture(url)
    img = best_image(cap)
    print(f"  Title: {cap['title']}")
    print(f"  Best image: {img}")
    print(f"  Quality: {assess_image_quality(cap['image_url'])} (original)")

    result = {"url": url, **cap, "best_image": img}

    if analyze and img:
        print("  Analyzing with Gemini Vision...")
        analysis = analyze_with_vision(img)
        result["analysis"] = analysis
        print(f"  Analysis: {json.dumps(analysis, indent=2)}")

    if notion_page_id and img:
        note = ""
        if "analysis" in result:
            a = result["analysis"]
            note = a.get("taste_note", json.dumps(a)[:200])
        ok = set_notion_cover(notion_page_id, img, note)
        print(f"  Notion cover set: {ok}")
        result["notion_updated"] = ok

    if output_dir:
        os.makedirs(output_dir, exist_ok=True)
        slug = re.sub(r'[^a-z0-9]+', '-', url.lower().split("//")[-1])[:60]
        with open(os.path.join(output_dir, f"{slug}.json"), "w") as f:
            json.dump(result, f, indent=2)

    return result


def main():
    p = argparse.ArgumentParser(description="Visual Intelligence Capture")
    p.add_argument("url", nargs="?", help="URL to capture")
    p.add_argument("--notion-page-id", help="Notion page ID to update cover")
    p.add_argument("--analyze", action="store_true", help="Run Gemini Vision analysis")
    p.add_argument("--output-dir", help="Save results as JSON files")
    p.add_argument("--batch", help="File with URLs (one per line)")
    args = p.parse_args()

    if args.batch:
        urls = [l.strip() for l in open(args.batch) if l.strip() and not l.startswith("#")]
        for i, url in enumerate(urls):
            if i > 0:
                time.sleep(2)
            process_url(url, analyze=args.analyze, output_dir=args.output_dir)
    elif args.url:
        r = process_url(args.url, args.notion_page_id, args.analyze, args.output_dir)
        print(json.dumps(r, indent=2))
    else:
        p.print_help()


if __name__ == "__main__":
    main()