#!/usr/bin/env python3
"""Local-only reference collector — licence-free, gitignored research material.

Downloads images from DuckDuckGo image search into `tmp/research/<slug>/` for
development reference only. This tool deliberately applies **no licence gate**:
everything it fetches stays under the gitignored `tmp/` tree and must never be
committed, shipped, or promoted into `docs/screens/refs/` without passing the
normal sidecar licence + privacy gate. The owner authorised this local-only
override on 2026-11-03 (see TODO.md §3D model re-evaluation → A).

    python3 tools/refcollect.py search "El Capitan Entertainment Centre facade" \
        --slug elcapitan --limit 20
    python3 tools/refcollect.py sheet --slug elcapitan

Requires the `ddgs` package (`pip install ddgs`); the legacy
`duckduckgo_search` package still works as a fallback.

Every outbound request goes through a token bucket and a polite retry routine:
DuckDuckGo and the image hosts rate-limit aggressive clients, so a large batch
is deliberately slow. Tune the pace with `--rate` (requests per second) and
`--burst`.
"""

from __future__ import annotations

import argparse
import datetime
import email.utils
import gzip
import json
import random
import re
import socket
import ssl
import sys
import threading
import time
import urllib.error
import urllib.request
import zlib
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
RESEARCH_DIR = ROOT / "tmp" / "research"

# --- politeness ------------------------------------------------------------

# A current desktop Firefox on Linux. Kept in one place so the search client and
# the downloader present the same, consistent identity.
USER_AGENT = ("Mozilla/5.0 (X11; Linux x86_64; rv:159.0) "
              "Gecko/20100101 Firefox/159.0")

# A coherent Firefox header set. Servers increasingly look past the user agent;
# a plain UA with no matching Accept/Accept-Language/Sec-Fetch headers is itself
# a tell, so send the whole set the browser would send.
#
# Accept-Encoding deliberately stops at gzip and deflate: urllib cannot decode
# brotli or zstd, and advertising an encoding we cannot read would corrupt every
# download. `_decode_body` undoes the two encodings we do advertise.
BROWSER_HEADERS = {
    "User-Agent": USER_AGENT,
    "Accept": ("image/avif,image/webp,image/png,image/svg+xml,image/*;q=0.8,"
               "*/*;q=0.5"),
    "Accept-Language": "en-US,en;q=0.9,de-CH;q=0.8,de;q=0.7",
    "Accept-Encoding": "gzip, deflate",
    "DNT": "1",
    "Sec-Fetch-Dest": "image",
    "Sec-Fetch-Mode": "no-cors",
    "Sec-Fetch-Site": "cross-site",
    "Sec-GPC": "1",
    "Referer": "https://duckduckgo.com/",
    "Upgrade-Insecure-Requests": "1",
}

DEFAULT_RATE = 0.5        # requests per second (one every two seconds)
DEFAULT_BURST = 3.0       # instant burst a cold start may spend
BACKOFF_BASE_S = 4.0      # first exponential-backoff ceiling
BACKOFF_CAP_S = 120.0     # never wait longer than two minutes
MAX_ATTEMPTS = 6
REQUEST_TIMEOUT_S = 60.0
RETRY_STATUS = frozenset(
    {408, 425, 429, 500, 502, 503, 504, 520, 521, 522, 524})


class TokenBucket:
    """Token bucket that paces outbound requests instead of fixed sleeps.

    Tokens accrue at ``rate`` per second up to ``capacity``. Each request spends
    one token; when the bucket is empty the caller sleeps only as long as it
    takes for the next token to accrue. A burst allowance lets a cold start fire
    a few requests back to back, which is what a human opening a page looks
    like, while the sustained rate stays gentle.
    """

    def __init__(self, rate: float, capacity: float | None = None) -> None:
        if rate <= 0:
            raise ValueError("rate must be positive")
        self.rate = float(rate)
        self.capacity = float(
            capacity if capacity is not None else max(1.0, rate))
        if self.capacity < 1.0:
            raise ValueError("capacity must hold at least one token")
        self._tokens = self.capacity
        self._stamp = time.monotonic()
        self._lock = threading.Lock()

    def _refill(self) -> None:
        now = time.monotonic()
        self._tokens = min(self.capacity,
                           self._tokens + (now - self._stamp) * self.rate)
        self._stamp = now

    @property
    def tokens(self) -> float:
        with self._lock:
            self._refill()
            return self._tokens

    def acquire(self, tokens: float = 1.0) -> float:
        """Block until ``tokens`` are available and return the seconds waited."""
        if tokens > self.capacity:
            raise ValueError("requested tokens exceed bucket capacity")
        waited = 0.0
        while True:
            with self._lock:
                self._refill()
                if self._tokens >= tokens:
                    self._tokens -= tokens
                    return waited
                delay = (tokens - self._tokens) / self.rate
            time.sleep(delay)
            waited += delay


def _parse_retry_after(value: str | None) -> float | None:
    """Read a Retry-After header (delta seconds or HTTP date) as seconds."""
    if not value:
        return None
    value = value.strip()
    if value.isdigit():
        return float(value)
    try:
        when = email.utils.parsedate_to_datetime(value)
    except (TypeError, ValueError):
        return None
    if when is None:
        return None
    if when.tzinfo is None:
        when = when.replace(tzinfo=datetime.timezone.utc)
    delta = (when - datetime.datetime.now(datetime.timezone.utc)).total_seconds()
    return max(0.0, delta)


def _backoff_delay(attempt: int, retry_after: float | None = None) -> float:
    """Full-jitter exponential backoff, capped, deferring to Retry-After.

    ``attempt`` starts at 1. When the host told us how long to wait, that wins,
    so we never come back sooner than asked. Otherwise pick uniformly from the
    top half of the exponential window: enough spread to avoid a thundering herd
    of retries, without doubling the total wait time.
    """
    if retry_after is not None:
        return min(BACKOFF_CAP_S, max(1.0, retry_after))
    ceiling = min(BACKOFF_CAP_S, BACKOFF_BASE_S * 2 ** (attempt - 1))
    return random.uniform(ceiling / 2.0, ceiling)


def _is_permanent_network_error(exc: BaseException) -> bool:
    """True when retrying cannot help: bad DNS or an untrusted certificate."""
    reason = getattr(exc, "reason", None) or exc
    return isinstance(reason, (socket.gaierror, ssl.SSLCertVerificationError))


def _decode_body(body: bytes, content_encoding: str | None) -> bytes:
    """Undo a Content-Encoding we advertised in BROWSER_HEADERS."""
    encoding = (content_encoding or "").lower().strip()
    if encoding in ("gzip", "x-gzip"):
        return gzip.decompress(body)
    if encoding == "deflate":
        try:
            return zlib.decompress(body)
        except zlib.error:
            return zlib.decompress(body, -zlib.MAX_WBITS)  # raw deflate
    return body


def polite_get(url: str, *, bucket: TokenBucket, headers: dict | None = None,
               attempts: int = MAX_ATTEMPTS, timeout: float = REQUEST_TIMEOUT_S,
               label: str = "request") -> tuple[bytes, object]:
    """Fetch ``url`` politely and return ``(body, response headers)``.

    Each attempt takes a token from ``bucket`` first. Transient failures are
    retried with jittered exponential backoff, honouring Retry-After; permanent
    failures raise immediately so a dead link does not burn the whole budget.
    """
    request_headers = dict(BROWSER_HEADERS)
    if headers:
        request_headers.update(headers)

    last_exc: BaseException | None = None
    for attempt in range(1, attempts + 1):
        bucket.acquire()
        retry_after = None
        try:
            request = urllib.request.Request(url, headers=request_headers)
            with urllib.request.urlopen(request, timeout=timeout) as response:
                body = _decode_body(response.read(),
                                    response.headers.get("Content-Encoding"))
                return body, response.headers
        except urllib.error.HTTPError as exc:
            if exc.code not in RETRY_STATUS:
                raise
            last_exc = exc
            retry_after = _parse_retry_after(exc.headers.get("Retry-After"))
            reason = f"HTTP {exc.code} {exc.reason or ''}".strip()
        except (urllib.error.URLError, OSError, ValueError) as exc:
            if _is_permanent_network_error(exc):
                raise
            last_exc = exc
            reason = exc.__class__.__name__
        if attempt == attempts:
            break
        delay = _backoff_delay(attempt, retry_after)
        print(f"  {label}: {reason}; politely waiting {delay:.1f}s "
              f"before attempt {attempt + 1}/{attempts}",
              file=sys.stderr, flush=True)
        time.sleep(delay)
    raise last_exc if last_exc is not None else RuntimeError(
        f"could not fetch {url}")


def _rel(path: Path) -> str:
    return str(path.relative_to(ROOT)) if path.is_relative_to(ROOT) else str(path)


def _slug(text: str) -> str:
    return re.sub(r"[^a-z0-9]+", "-", text.strip().lower()).strip("-") or "block"


def _search_images(query: str, limit: int, bucket: TokenBucket) -> list[dict]:
    """Return DDG image results as small dicts; polite retry on rate limits."""
    try:
        from ddgs import DDGS  # current package name
    except ImportError:
        # Legacy `duckduckgo_search` still works but warns it was renamed to
        # `ddgs`; silence that warning so it does not leak into the run log.
        import warnings
        try:
            with warnings.catch_warnings():
                warnings.filterwarnings(
                    "ignore", category=RuntimeWarning,
                    message=r".*renamed to .ddgs.*")
                from duckduckgo_search import DDGS  # legacy import name
        except ImportError as exc:
            sys.exit("error: install the DuckDuckGo search package first: "
                     "pip install ddgs")
    last_exc = None
    for attempt in range(1, MAX_ATTEMPTS + 1):
        bucket.acquire()
        try:
            with DDGS() as ddgs:
                rows = ddgs.images(query, max_results=limit)
            return [{
                "title": (r.get("title") or "")[:200],
                "image": r.get("image") or "",
                "thumbnail": r.get("thumbnail") or "",
                "source": r.get("url") or "",
                "width": r.get("width") or 0,
                "height": r.get("height") or 0,
            } for r in rows if r.get("image")]
        except Exception as exc:  # noqa: BLE001 - package raises several types
            last_exc = exc
            if attempt == MAX_ATTEMPTS:
                break
            delay = _backoff_delay(attempt)
            print(f"  DDG search failed ({exc.__class__.__name__}); politely "
                  f"waiting {delay:.1f}s before attempt {attempt + 1}/"
                  f"{MAX_ATTEMPTS}", file=sys.stderr, flush=True)
            time.sleep(delay)
    sys.exit(f"error: DuckDuckGo image search kept failing: {last_exc}")


def _download(url: str, path: Path, bucket: TokenBucket) -> None:
    """Download one image through the polite, rate-limited fetch routine."""
    body, _ = polite_get(url, bucket=bucket, label=f"download {path.name}")
    path.write_bytes(body)


def _manifest_rows(rows: list[dict]) -> list[dict]:
    out = []
    for row in rows:
        out.append({**row, "local_only": True, "licence": "unchecked-local"})
    return out


def cmd_search(args) -> int:
    slug = _slug(args.slug or args.query)
    out_dir = RESEARCH_DIR / slug
    out_dir.mkdir(parents=True, exist_ok=True)
    bucket = TokenBucket(args.rate, args.burst)
    if getattr(args, "user_agent", ""):
        BROWSER_HEADERS["User-Agent"] = args.user_agent
    print(f"searching DDG: {args.query!r} (limit {args.limit}, "
          f"{args.rate:g} req/s, burst {args.burst:g})", flush=True)
    rows = _search_images(args.query, args.limit, bucket)
    manifest = []
    for index, row in enumerate(rows, start=1):
        stem = f"{index:02d}"
        suffix = Path(row["image"].split("?")[0]).suffix.lower() or ".jpg"
        if suffix not in (".jpg", ".jpeg", ".png", ".webp", ".gif"):
            suffix = ".jpg"
        path = out_dir / f"{stem}{suffix}"
        if path.exists() and not args.force:
            print(f"  skip existing {path.name}")
        else:
            try:
                _download(row["image"], path, bucket)
                print(f"  {path.name} <- {row['image'][:70]}")
            except Exception as exc:  # noqa: BLE001 - some hosts block hotlinking
                print(
                    f"  FAILED {path.name}: {exc.__class__.__name__}", file=sys.stderr)
                continue
        manifest.append({**row, "file": path.name})
    (out_dir / "manifest.json").write_text(
        json.dumps(_manifest_rows(manifest), indent=2) + "\n")
    print(f"manifest -> {out_dir / 'manifest.json'} ({len(manifest)} entries)")
    return 0


def cmd_sheet(args) -> int:
    """Compose a local contact sheet (never committed) from a manifest folder."""
    from PIL import Image
    slug = _slug(args.slug)
    out_dir = RESEARCH_DIR / slug
    manifest_path = out_dir / "manifest.json"
    if not manifest_path.is_file():
        sys.exit(f"error: no manifest at {manifest_path}; run `search` first")
    rows = json.loads(manifest_path.read_text())
    images = []
    for row in rows:
        path = out_dir / row["file"]
        if not path.is_file():
            continue
        try:
            with Image.open(path) as img:
                img.load()
            images.append(img.convert("RGB"))
        except Exception as exc:  # noqa: BLE001 - AVIF/misnamed or corrupt frames
            print(f"  skip unreadable {path.name}: {exc.__class__.__name__}",
                  file=sys.stderr)
    if not images:
        sys.exit("error: no readable frames in this folder (convert AVIF/webp "
                 "files with ffmpeg first)")
    cell = 320
    cols = 4
    rows_count = (len(images) + cols - 1) // cols
    sheet = Image.new("RGB", (cols * cell, rows_count * cell), (18, 18, 18))
    for i, img in enumerate(images):
        thumb = img.copy()
        thumb.thumbnail((cell, cell), Image.Resampling.LANCZOS)
        x = (i % cols) * cell
        y = (i // cols) * cell
        sheet.paste(thumb, (x, y))
    out = out_dir / "sheet.jpg"
    sheet.save(out, quality=82)
    print(f"wrote {_rel(out)} {sheet.size[0]}x{sheet.size[1]} "
          f"({len(images)} frames)")
    return 0


def build_parser():
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = parser.add_subparsers(dest="cmd", required=True)
    search = sub.add_parser(
        "search", help="DDG image search into tmp/research/<slug>/")
    search.add_argument("query")
    search.add_argument("--slug", default="",
                        help="folder slug (default: from query)")
    search.add_argument("--limit", type=int, default=20)
    search.add_argument("--force", action="store_true")
    search.add_argument("--rate", type=float, default=DEFAULT_RATE,
                        help="requests per second, enforced by a token bucket "
                             f"(default {DEFAULT_RATE:g})")
    search.add_argument("--burst", type=float, default=DEFAULT_BURST,
                        help="token-bucket capacity, i.e. allowed instant burst "
                             f"(default {DEFAULT_BURST:g})")
    search.add_argument("--user-agent", default="",
                        help="override the default Firefox user agent string")
    search.set_defaults(handler=cmd_search)
    sheet = sub.add_parser("sheet", help="compose a local contact sheet")
    sheet.add_argument("--slug", required=True)
    sheet.set_defaults(handler=cmd_sheet)
    return parser


def main(argv=None) -> int:
    args = build_parser().parse_args(argv)
    return args.handler(args)


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