queue_both_tabs.py

"""One-shot: queue everything from both tabs files that isn't already in the DB.

Differs from dedup_and_queue_tabs.py: does NOT use tabs-organized.txt as a "known"
filter — only filters against the actual installed-content signals in postgres
(have_set, source_url, creators table, tags table for nhentai).
"""

from __future__ import annotations
import argparse, json, sys, time, urllib.error, urllib.request
from collections import Counter
from pathlib import Path

import psycopg2
from gap_analysis import PG_DSN
from gap_analysis.core import load_have_set, normalize_slug
from gap_analysis.known import _load_from_postgres  # DB-only nhentai known creators
from ingest_router import detect_route

INGEST = "http://127.0.0.1:8090/api/ingest/url"


def parse_tabs(path: Path) -> list[str]:
    out = []
    for raw in path.read_text(encoding="utf-8", errors="replace").splitlines():
        s = raw.strip()
        if s.lower().startswith(("http://", "https://")):
            out.append(s)
    return out


def load_source_urls(dsn):
    conn = psycopg2.connect(dsn)
    try:
        with conn.cursor() as cur:
            cur.execute("SELECT source_url FROM posts WHERE source_url IS NOT NULL")
            return {r[0] for r in cur if r[0]}
    finally:
        conn.close()


def load_creator_slugs(dsn):
    from collections import defaultdict
    out = defaultdict(set)
    conn = psycopg2.connect(dsn)
    try:
        with conn.cursor() as cur:
            cur.execute("SELECT service, username, user_id FROM creators WHERE username IS NOT NULL OR user_id IS NOT NULL")
            for service, username, user_id in cur:
                if username: out[service].add(normalize_slug(username))
                if user_id: out[service].add(normalize_slug(user_id))
    finally:
        conn.close()
    return out


def is_dup(url, route, nh_have, nh_known_db, source_urls, creators_by_service):
    if url in source_urls:
        return "source_url"
    if route is None:
        return None
    if route.handler == "nhentai_gallery" and route.content_id:
        try:
            if int(route.content_id) in nh_have: return "nh-gallery"
        except ValueError:
            pass
    elif route.handler in ("nhentai_artist", "nhentai_group") and route.content_id:
        kind = "artist" if route.handler == "nhentai_artist" else "group"
        slug = normalize_slug(route.content_id)
        if (kind, slug) in nh_known_db: return f"nh-{kind}-db"
    elif route.handler == "pornhub_creator" and route.content_id:
        slug = normalize_slug(route.content_id)
        if slug in creators_by_service.get("pornhub", set()): return "ph-creator"
    elif route.handler == "coomer_creator" and route.content_id:
        slug = normalize_slug(route.content_id)
        for svc in ("coomer","kemono","onlyfans","fansly"):
            if slug in creators_by_service.get(svc, set()): return f"{svc}-creator"
    elif route.handler == "hentairead_artist" and route.content_id:
        slug = normalize_slug(route.content_id)
        if slug in creators_by_service.get("hentairead", set()): return "hr-creator"
    return None


def post_url(url, timeout=15):
    body = json.dumps({"url": url}).encode()
    req = urllib.request.Request(INGEST, data=body, method="POST",
                                  headers={"Content-Type":"application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--tabs", nargs="+", required=True)
    ap.add_argument("--queue", action="store_true")
    ap.add_argument("--delay", type=float, default=0.02)
    ap.add_argument("--out", type=str, default=None,
                    help="write the not-yet-downloaded (deduped) URL list to this path")
    args = ap.parse_args()

    # 1. Collect URLs, dedupe by URL string THEN by (handler, content_id)
    #    so /g/123/ and /g/123/57/ (same gallery) collapse to one task.
    raw = set()
    for p in args.tabs:
        raw.update(parse_tabs(Path(p)))
    seen_content = set()
    urls = []
    collapsed = 0
    for u in sorted(raw):
        r = detect_route(u)
        key = (r.handler, r.content_id) if (r and r.content_id) else ("url", u)
        if key in seen_content:
            collapsed += 1
            continue
        seen_content.add(key)
        urls.append(u)
    print(f"[1] {len(raw)} URL-unique → {len(urls)} content-unique "
          f"(-{collapsed} same-content variants)", file=sys.stderr)

    # 2. DB-only dedup sets
    print("[2] loading DB sets...", file=sys.stderr)
    nh_have = load_have_set(PG_DSN)
    print(f"      nhentai have-set: {len(nh_have)}", file=sys.stderr)
    nh_known_db = _load_from_postgres(PG_DSN)  # nhentai tags table ONLY
    print(f"      nhentai DB-known creators (tags table): {len(nh_known_db)}", file=sys.stderr)
    source_urls = load_source_urls(PG_DSN)
    print(f"      posts.source_url: {len(source_urls)}", file=sys.stderr)
    creators_by_service = load_creator_slugs(PG_DSN)
    print(f"      creators across {len(creators_by_service)} services", file=sys.stderr)

    # 3. Filter
    keep, drop = [], []
    drop_reasons = Counter()
    handlers = Counter()
    for u in urls:
        route = detect_route(u)
        h = route.handler if route else "rejected"
        handlers[h] += 1
        reason = is_dup(u, route, nh_have, nh_known_db, source_urls, creators_by_service)
        if reason:
            drop.append((u, reason)); drop_reasons[reason] += 1
        else:
            keep.append((u, h))
    print(f"[3] kept={len(keep)} dropped={len(drop)}", file=sys.stderr)
    print("    handlers:", file=sys.stderr)
    for h,n in handlers.most_common(): print(f"      {h:22s} {n}", file=sys.stderr)
    print("    drop reasons:", file=sys.stderr)
    for r,n in drop_reasons.most_common(): print(f"      {r:22s} {n}", file=sys.stderr)

    if args.out:
        Path(args.out).write_text(
            "\n".join(u for u, _ in keep) + "\n", encoding="utf-8")
        print(f"[out] wrote {len(keep)} not-yet-downloaded URLs to {args.out}",
              file=sys.stderr)

    if not args.queue:
        print("[4] dry-run; pass --queue to POST", file=sys.stderr)
        return 0

    # 4. POST
    ok = dup = fail = 0
    for i,(u,h) in enumerate(keep, 1):
        try:
            r = post_url(u)
            if r.get("duplicate"): dup += 1
            else: ok += 1
        except Exception as e:
            fail += 1
            if fail <= 3: print(f"  fail {u}: {e}", file=sys.stderr)
        if i % 500 == 0:
            print(f"  [{i}/{len(keep)}] ok={ok} dup={dup} fail={fail}", file=sys.stderr)
        if args.delay > 0: time.sleep(args.delay)
    print(f"\n[4] DONE: ok={ok} dup={dup} fail={fail}", file=sys.stderr)
    return 0 if fail == 0 else 1


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