#!/usr/bin/env python3
"""
Search Pinterest (via pinch.py), download images, make clean transparent cutouts.

Pipeline per image:
  pinch download -> birefnet cutout -> (scene-fail? retry heavier model)
  -> defringe edges -> clean alpha -> trim margins -> resize for web -> save

Reuses the search/download functions from pinch.py (same folder).

Usage:
    pip install requests rembg pillow onnxruntime numpy scipy
    python scripts/fetch_qts.py            # fetch & append to the library
    python scripts/fetch_qts.py --recut    # re-clean images already in the folder

It asks:
    image to find:       (e.g. "hyuuga_hinata")
    how many to find:
    search hint:         (optional; blank = search exactly what you type)

Output: transparent PNGs with random names in  public/assets/qts/ .
New runs APPEND (never overwrite), skip duplicates, and it rewrites the
QTS array in public/assets/app.js automatically.
"""

import argparse
import hashlib
import json
import re
import secrets
import sys
import time
from html import unescape
from pathlib import Path
from io import BytesIO

# pinch.py lives next to this file
sys.path.insert(0, str(Path(__file__).resolve().parent))
import pinch  # noqa: E402  (user's Pinterest downloader)

try:
    import numpy as np
    from rembg import remove, new_session
    from PIL import Image
    HAVE_REMBG = True
except ImportError:
    HAVE_REMBG = False

HERE = Path(__file__).resolve().parent
OUT_DIR = HERE.parent / "public" / "assets" / "qts"
APP_JS = HERE.parent / "public" / "assets" / "app.js"
HASHES_FILE = OUT_DIR / ".hashes.txt"  # remembers downloaded images to skip dupes
MAX_COUNT = 2000

# ============================ quality settings ============================

# Where to search:
#   "auto"   booru first (reliable anime tags), fall back to pinterest/bing
#   "booru"  booru only (anime; never blocks)
#   "scrape" pinterest native + bing only (can get rate-limited / return junk)
SOURCE = "auto"

# Primary cutout model (first use downloads it once):
#   "birefnet-general"  SOTA matting, separates subject from busy backgrounds.
#   "birefnet-massive"  biggest/best, slower — used as the scene-fail fallback.
#   "isnet-anime"       fast, good only on already-clean 2D renders.
MODEL = "birefnet-general"

# Quality gate (the biggest quality win): only KEEP cutouts whose kept area is in
# a healthy band, measured before trim. Above GOOD_MAX = the model failed to cut
# (whole scene kept); below GOOD_MIN = it wiped the subject. Bad cuts are skipped
# instead of saved, so only clean cutouts land in qts.
GOOD_MIN = 0.05
GOOD_MAX = 0.90
# Fetch extra candidates so skipped (bad) cuts still leave enough good ones.
OVERFETCH = 3

# Optional heavier retry for borderline cuts. birefnet-massive tends to over-cut
# anime, so this is OFF by default. Set to "birefnet-massive" to experiment.
FALLBACK_MODEL = None
SCENE_FAIL_RATIO = 0.90   # >this opaque => try fallback (if enabled)
FALLBACK_MIN_KEEP = 0.05  # reject a fallback that keeps less than this (over-cut)

# Search hint is asked for each run (blank = search exactly what you type).
# This is just the suggestion shown at the prompt.
HINT_SUGGESTION = "render transparent png"

# Keep only the largest connected subject (drops floating extras, e.g. a second
# object on a product photo). Needs scipy. Set False to keep every object.
KEEP_LARGEST = True

# Edge cleanup
DEFRINGE = True           # erode the 1px contaminated rim that causes color halos
CLEAN_ALPHA = True        # snap near-0 / near-255 alpha so no faint background ghost
TRIM = True               # crop away fully transparent margins
MAX_DIM = 1200            # downscale so qts load fast on the page (0 = keep original)

# Alpha matting (only helps hard-masked models like isnet/u2net; birefnet feathers
# edges itself). Needs:  pip install pymatting
ALPHA_MATTING = False
AM_FG, AM_BG, AM_ERODE = 240, 15, 12

try:
    import pymatting  # noqa: F401
    HAVE_PYMATTING = True
except ImportError:
    HAVE_PYMATTING = False

try:
    from scipy import ndimage  # for keep-largest-subject
    HAVE_SCIPY = True
except ImportError:
    HAVE_SCIPY = False

# ==========================================================================


IMG_SUFFIX = (".jpg", ".jpeg", ".png", ".webp")

# ---- booru source (Safebooru: SFW anime, JSON API, no key, no bot-blocking) ----
BOORU_API = "https://safebooru.org/index.php"
BOORU_AC = "https://safebooru.org/autocomplete.php"
BOORU_PER = 100


def booru_suggest(session, q):
    """Autocomplete real booru tags for `q` -> [(tag, post_count), ...]."""
    try:
        items = session.get(BOORU_AC, params={"q": q}, timeout=pinch.TIMEOUT).json()
    except Exception:
        return []
    out = []
    for it in items:
        tag = it.get("value", "")
        m = re.search(r"\((\d+)\)", it.get("label", ""))
        if tag:
            out.append((tag, int(m.group(1)) if m else 0))
    return out


def booru_search(session, tags, count):
    """Collect up to `count` image URLs for `tags` from Safebooru."""
    urls, pid = [], 0
    while len(urls) < count:
        try:
            posts = session.get(BOORU_API, params={
                "page": "dapi", "s": "post", "q": "index", "json": "1",
                "tags": tags, "limit": BOORU_PER, "pid": pid,
            }, timeout=pinch.TIMEOUT).json()
        except Exception:
            break
        if not posts:
            break
        for p in posts:
            u = p.get("file_url")
            if not u and p.get("directory") is not None and p.get("image"):
                u = f"https://safebooru.org/images/{p['directory']}/{p['image']}"
            if u and u.lower().endswith(IMG_SUFFIX):
                urls.append(u)
                if len(urls) >= count:
                    break
        pid += 1
    return urls


# Generic "anime girl" style queries -> the real booru tag for them.
BOORU_GENERIC = {
    "waifu": "1girl solo", "waifus": "1girl solo",
    "anime girl": "1girl solo", "anime girls": "1girl solo",
    "anime": "1girl solo", "girl": "1girl solo", "girls": "1girl solo",
}


def booru_resolve(session, raw):
    """Map whatever the user typed to a real booru tag (with a 'did you mean' picker)."""
    if raw.lower() in BOORU_GENERIC:
        return BOORU_GENERIC[raw.lower()]
    if booru_search(session, raw, 1):
        return raw
    joined = raw.replace(" ", "_")
    if joined != raw and booru_search(session, joined, 1):
        return joined
    seen, cands = set(), []
    for q in [joined] + raw.split():
        for tag, c in booru_suggest(session, q):
            if tag not in seen:
                seen.add(tag)
                cands.append((tag, c))
    words = [w for w in raw.lower().split() if w]
    cands.sort(key=lambda it: (sum(1 for w in words if w in it[0].lower()), it[1]), reverse=True)
    cands = cands[:12]
    if not cands:
        return None
    print(f"\nno exact booru tag for '{raw}'. did you mean:")
    for i, (tag, c) in enumerate(cands, 1):
        print(f"  {i:2}. {tag}  ({c})")
    choice = input("pick a number (enter to skip booru): ").strip()
    if not choice.isdigit() or not (1 <= int(choice) <= len(cands)):
        return None
    return cands[int(choice) - 1][0]


def bing_any(session, query, first):
    """Bing image results for `query` — ANY image host (not just pinimg)."""
    try:
        r = session.get(
            pinch.BING_IMAGES,
            params={"q": query, "first": first, "count": pinch.BING_PAGE},
            timeout=pinch.TIMEOUT,
        )
        if r.status_code != 200:
            return []
    except Exception:
        return []
    out = []
    for raw in re.findall(r'm="(\{[^"]+\})"', r.text):
        try:
            item = json.loads(unescape(raw))
        except Exception:
            continue
        u = (item.get("murl") or "").split("?")[0]
        if u.startswith("http") and u.lower().endswith(IMG_SUFFIX):
            out.append(u)
    return out


def collect_images(session, query, count):
    """Booru first (reliable), then Pinterest native, then a permissive Bing search."""
    # --- booru (Safebooru) ---
    if SOURCE in ("auto", "booru"):
        tag = booru_resolve(session, query)
        if tag:
            urls = booru_search(session, tag, count)
            if urls:
                print(f"  booru tag '{tag}' — {len(urls)} result(s)")
                return urls[:count]
        if SOURCE == "booru":
            return []  # booru-only mode, nothing found

    # --- scrape fallback ---
    urls = []
    try:
        urls = pinch.collect_native(session, query, count, quiet=False)
    except Exception:
        pass
    if len(urls) >= count:
        return urls[:count]

    print("  (pinterest short/blocked — using general image search)")
    seen = set(urls)
    # Page the EXACT query deep — never append noise words, they wreck relevance.
    first, stall = 1, 0
    while len(urls) < count and stall < 6:
        fresh = [u for u in bing_any(session, query, first) if u not in seen]
        if fresh:
            stall = 0
            seen.update(fresh)
            urls.extend(fresh)
        else:
            stall += 1
        print(f"  [search] '{query}' {len(urls)}/{count}", file=sys.stderr)
        first += pinch.BING_PAGE
        time.sleep(pinch.SLEEP)
    return urls[:count]


def get_image_bytes(session, url):
    """Download one image, trying pinch's url size variants."""
    for variant in pinch.url_variants(url):
        try:
            got = pinch.fetch_bytes(session, variant)
        except Exception:
            got = None
        if got:
            return got[0]  # (bytes, ext) -> bytes
    return None


def opaque_ratio(img):
    """Fraction of pixels that are mostly opaque — a cheap 'did the cut work?' signal."""
    a = np.asarray(img)[:, :, 3]
    return float((a > 128).mean())


def erode_alpha(a):
    """Shrink the alpha mask by ~1px (4-neighbour min) to drop the fringe rim."""
    out = a.copy()
    out[1:, :] = np.minimum(out[1:, :], a[:-1, :])
    out[:-1, :] = np.minimum(out[:-1, :], a[1:, :])
    out[:, 1:] = np.minimum(out[:, 1:], a[:, :-1])
    out[:, :-1] = np.minimum(out[:, :-1], a[:, 1:])
    return out


def refine(img):
    """Defringe + clean the alpha channel for a crisp, halo-free edge."""
    arr = np.array(img).astype(np.uint8)
    a = arr[:, :, 3]
    if DEFRINGE:
        a = erode_alpha(a)
    if CLEAN_ALPHA:
        a = np.where(a < 12, 0, a)
        a = np.where(a > 243, 255, a).astype(np.uint8)
    arr[:, :, 3] = a
    return Image.fromarray(arr, "RGBA")


def keep_largest(img):
    """Keep only the biggest connected subject; drop smaller floating objects."""
    if not (KEEP_LARGEST and HAVE_SCIPY):
        return img
    arr = np.array(img)
    a = arr[:, :, 3]
    mask = a > 10
    labels, n = ndimage.label(mask)
    if n <= 1:
        return img
    counts = np.bincount(labels.ravel())
    counts[0] = 0  # ignore the transparent background label
    keep = counts.argmax()
    arr[:, :, 3] = np.where(labels == keep, a, 0).astype(np.uint8)
    return Image.fromarray(arr, "RGBA")


def trim(img):
    """Crop away fully transparent borders."""
    bbox = img.getbbox()  # bounds of non-zero (non-transparent) region
    return img.crop(bbox) if bbox else img


def random_name():
    """A unique random filename for a new cutout."""
    while True:
        name = secrets.token_hex(8) + ".png"
        if not (OUT_DIR / name).exists():
            return name


def cap(img):
    """Downscale so the longest side is <= MAX_DIM."""
    if not MAX_DIM:
        return img
    w, h = img.size
    if max(w, h) <= MAX_DIM:
        return img
    scale = MAX_DIM / max(w, h)
    return img.resize((round(w * scale), round(h * scale)), Image.LANCZOS)


def cut_out(img, primary, fallback):
    """Cut the background, retrying with the heavier model if the first cut fails."""
    kwargs = dict(
        post_process_mask=True,
        alpha_matting=(ALPHA_MATTING and HAVE_PYMATTING),
        alpha_matting_foreground_threshold=AM_FG,
        alpha_matting_background_threshold=AM_BG,
        alpha_matting_erode_size=AM_ERODE,
    )
    cut = remove(img, session=primary, **kwargs)
    r = opaque_ratio(cut)
    if fallback is not None and r > SCENE_FAIL_RATIO:
        alt = remove(img, session=fallback, **kwargs)
        ra = opaque_ratio(alt)
        # only accept the heavier cut if it removed more AND kept a real subject
        if FALLBACK_MIN_KEEP < ra < r:
            print(f"    fallback kept {ra:.0%} (was {r:.0%})")
            cut, r = alt, ra
    return cut, r  # r = kept-area ratio before trim, for the quality gate


def library_pngs():
    """All cutouts in the folder, oldest first (stable cycle order)."""
    return sorted(OUT_DIR.glob("*.png"), key=lambda p: p.stat().st_mtime)


def load_hashes():
    try:
        return set(HASHES_FILE.read_text(encoding="utf-8").split())
    except Exception:
        return set()


def remember_hash(h):
    try:
        with HASHES_FILE.open("a", encoding="utf-8") as fh:
            fh.write(h + "\n")
    except Exception:
        pass


def write_app_js():
    """Rewrite the QTS = [...] array in app.js with EVERY cutout in the folder."""
    names = [p.name for p in library_pngs()]
    body = "".join(f"  'assets/qts/{n}',\n" for n in names)
    block = "const QTS = [\n" + body + "];"
    try:
        text = APP_JS.read_text(encoding="utf-8")
        new, count = re.subn(r"const QTS = \[.*?\];", block, text, count=1, flags=re.S)
        if count:
            APP_JS.write_text(new, encoding="utf-8")
            return len(names)
    except Exception as e:
        print(f"! could not update app.js automatically: {e}")
    return None


def make_models():
    """Load the cutout model session(s)."""
    primary = new_session(MODEL)
    fallback = new_session(FALLBACK_MODEL) if FALLBACK_MODEL else None
    print(f"cutout model: {MODEL}" + (f" (fallback: {FALLBACK_MODEL})" if fallback else ""))
    return primary, fallback


def clean_existing(primary, fallback, path):
    """Re-clean one library image in place. Returns True if it's a clean cutout now."""
    img = Image.open(path).convert("RGBA")
    if opaque_ratio(img) > GOOD_MAX:
        # looks un-cut (full scene still there) -> attempt a real cut
        cut, kept = cut_out(img, primary, fallback)
        if not (GOOD_MIN <= kept <= GOOD_MAX):
            return False  # still can't separate a subject
        img = cut
    # light cleanup for everything (good cutouts just get tidied, not re-modelled)
    img = keep_largest(img)
    img = refine(img)
    if TRIM:
        img = trim(img)
    img = cap(img)
    img.save(path)
    return True


def run_recut():
    """Re-run the cleanup pipeline over images already in the folder (no downloads)."""
    print("== qts recut (cleaning existing folder) ==")
    if not HAVE_REMBG:
        sys.exit("rembg/numpy/Pillow required for --recut.")
    pngs = library_pngs()
    if not pngs:
        sys.exit("No images in the folder to recut.")
    primary, fallback = make_models()
    print()
    bad = []
    for i, p in enumerate(pngs, start=1):
        print(f"[{i}/{len(pngs)}] {p.name}")
        try:
            if not clean_existing(primary, fallback, p):
                bad.append(p.name)
                print("    not cleanly transparent — left as-is")
        except Exception as e:
            print(f"    ! failed: {e}")
    total = write_app_js()
    print(f"\nRecut done. {len(pngs)} processed, {len(pngs) - len(bad)} clean.")
    if bad:
        print("Couldn't cut (still has a background): " + ", ".join(bad))
    if total is not None:
        print(f"app.js updated with all {total}. Hard-refresh the page.")


def run_fetch():
    print("== qts fetcher (pinterest) ==")
    query = input("image to find: ").strip().lstrip("")
    if not query:
        sys.exit("No search entered.")
    try:
        want = int(input("how many to find: ").strip())
    except ValueError:
        sys.exit("That wasn't a number.")
    want = max(1, min(want, MAX_COUNT))

    hint = input(f"search hint (optional, blank = exact) [{HINT_SUGGESTION}]: ").strip().lstrip("")

    OUT_DIR.mkdir(parents=True, exist_ok=True)

    # --- search Pinterest with pinch ---
    session = pinch.make_session()
    try:
        session.get(pinch.HOME_URL, timeout=pinch.TIMEOUT)
    except Exception:
        pass

    search = f"{query} {hint}".strip() if hint else query
    # Collecting URLs is cheap (no downloads), so pull a big pool and page deep —
    # this is what lets us skip past images already in the library to find NEW ones.
    want_candidates = want if not HAVE_REMBG else min(MAX_COUNT, max(want * OVERFETCH, 300))
    print(f"\nsearching for '{search}' ...")
    urls = collect_images(session, search, want_candidates)
    if not urls:
        sys.exit("No results found.")
    print(f"found {len(urls)} candidate(s).\n")

    # --- cutout models ---
    primary = fallback = None
    if HAVE_REMBG:
        primary, fallback = make_models()
        print()
    else:
        print("! rembg/numpy/Pillow not installed — saving WITHOUT transparency.")
        print("  pip install rembg pillow onnxruntime numpy\n")

    # --- download + cut out, APPENDING clean cutouts to the library ---
    # Existing images are kept; new ones get random names and skip dupes.
    start_count = len(library_pngs())
    if start_count == 0:
        # folder was cleared -> forget old dedup memory so we can re-pull
        HASHES_FILE.unlink(missing_ok=True)
    seen = load_hashes()
    added = dups = 0
    for i, url in enumerate(urls, start=1):
        if added >= want:
            break
        raw = get_image_bytes(session, url)
        if not raw:
            continue
        h = hashlib.md5(raw).hexdigest()
        if h in seen:
            dups += 1
            continue  # already in library — quietly skip
        print(f"[{added}/{want} new | library {start_count + added}] ({i}) {url}")
        out = OUT_DIR / random_name()
        try:
            if HAVE_REMBG:
                img = Image.open(BytesIO(raw)).convert("RGBA")
                cut, kept = cut_out(img, primary, fallback)
                if not (GOOD_MIN <= kept <= GOOD_MAX):
                    reason = "failed to cut" if kept > GOOD_MAX else "over-cut"
                    print(f"    skipped bad cutout ({reason}, kept {kept:.0%})")
                    continue
                cut = keep_largest(cut)
                cut = refine(cut)
                if TRIM:
                    cut = trim(cut)
                cut = cap(cut)
                cut.save(out)
            else:
                out.write_bytes(raw)
            seen.add(h)
            remember_hash(h)
            added += 1
        except Exception as e:
            print(f"    ! processing failed: {e}")

    total = write_app_js()
    print(f"\nDone. Added {added} new ({dups} duplicates skipped), "
          f"library now has {len(library_pngs())} in {OUT_DIR}")
    if added < want:
        print(f"Only found {added}/{want} new — this search is mostly exhausted. "
              f"Try a different or more specific query to get more variety.")
    if total is not None:
        print(f"app.js QTS array updated with all {total}. Just hard-refresh the page.")
    else:
        print("\n(couldn't auto-update app.js — the QTS array lists the folder's *.png files)")


def main():
    ap = argparse.ArgumentParser(description="fetch/clean transparent qts")
    ap.add_argument("--recut", action="store_true",
                    help="re-clean images already in the folder (no downloads)")
    args = ap.parse_args()
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    if args.recut:
        run_recut()
    else:
        run_fetch()


if __name__ == "__main__":
    main()