queue_organized_creators.py

"""Queue all creator URLs in tabs-organized.txt that aren't already finished.

"Finished" semantics per service:
  - nhentai_artist / nhentai_group : at least one post in DB tagged with this artist/group
    (uses gap_analysis.known._load_from_postgres — DB side only, NOT the tabs side,
    since tabs-organized.txt is the file we're processing — using its own contents
    as the exclude-set would zero out everything.)
  - pornhub_creator / hentairead_artist / coomer_creator / beeg_creator :
    a row exists in `creators` for that service+slug.

Additionally skip any URL whose task is currently QUEUED or DOWNLOADING on
the live server (so we don't pile duplicate work on the worker).

Non-creator URLs in the file (galleries, videos, etc.) are ignored — user
asked for creators specifically.

Run:
  ./.venv/bin/python queue_organized_creators.py                 # dry-run
  ./.venv/bin/python queue_organized_creators.py --queue         # POST surviving URLs
"""

from __future__ import annotations

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

import psycopg2

from gap_analysis import PG_DSN
from gap_analysis.core import normalize_slug
from gap_analysis.known import _load_from_postgres  # DB side only
from ingest_router import detect_route


TABS_FILE = Path(__file__).parent / "tabs-organized.txt"
INGEST_ENDPOINT = "http://127.0.0.1:8090/api/ingest/url"
TASKS_ENDPOINT = "http://127.0.0.1:8090/api/ingest/url/tasks"

# Handlers we treat as "creator-level" (i.e. fan-out content, worth re-queuing
# to catch new uploads). Galleries/videos are atomic and excluded.
CREATOR_HANDLERS = {
    "nhentai_artist",
    "nhentai_group",
    "pornhub_creator",
    "hentairead_artist",
    "coomer_creator",
    "beeg_creator",
}


def parse_urls(path: Path) -> list[str]:
    urls: list[str] = []
    for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
        s = line.strip()
        if not s or s.startswith("=") or not s.lower().startswith(("http://", "https://")):
            continue
        urls.append(s)
    return urls


def load_creator_slugs_by_service(pg_dsn: str) -> dict[str, set[str]]:
    out: dict[str, set[str]] = defaultdict(set)
    conn = psycopg2.connect(pg_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 load_inflight_urls() -> set[str]:
    """URLs currently queued or downloading on the live server."""
    try:
        req = urllib.request.Request(TASKS_ENDPOINT, headers={"Accept": "application/json"})
        with urllib.request.urlopen(req, timeout=15) as resp:
            tasks = json.loads(resp.read().decode("utf-8"))
        return {
            t.get("url")
            for t in tasks
            if t.get("status") in ("queued", "downloading") and t.get("url")
        }
    except Exception as e:  # noqa: BLE001
        print(f"WARN: could not load in-flight tasks: {e}", file=sys.stderr)
        return set()


def classify(
    url: str,
    nhentai_db_known: set[tuple[str, str]],
    creators_by_service: dict[str, set[str]],
    inflight: set[str],
) -> tuple[str, str]:
    """Return (decision, reason). decision in {keep, drop:<reason>, skip:not-creator}."""
    route = detect_route(url)
    if route is None or route.handler not in CREATOR_HANDLERS:
        return "skip", "non-creator"

    if url in inflight:
        return "drop", "in-flight"

    cid = (route.content_id or "").strip("/")
    slug = normalize_slug(cid)
    if not slug:
        return "skip", "no-slug"

    if route.handler == "nhentai_artist":
        if ("artist", slug) in nhentai_db_known:
            return "drop", "nhentai-tag-in-db"
    elif route.handler == "nhentai_group":
        if ("group", slug) in nhentai_db_known:
            return "drop", "nhentai-tag-in-db"
    elif route.handler == "pornhub_creator":
        if slug in creators_by_service.get("pornhub", set()):
            return "drop", "pornhub-creator-in-db"
    elif route.handler == "hentairead_artist":
        if slug in creators_by_service.get("hentairead", set()):
            return "drop", "hentairead-creator-in-db"
    elif route.handler == "coomer_creator":
        for svc in ("coomer", "kemono", "onlyfans", "fansly"):
            if slug in creators_by_service.get(svc, set()):
                return "drop", f"{svc}-creator-in-db"
    elif route.handler == "beeg_creator":
        if slug in creators_by_service.get("beeg", set()):
            return "drop", "beeg-creator-in-db"

    return "keep", route.handler


def post(url: str, timeout: float = 15.0) -> dict:
    body = json.dumps({"url": url}).encode("utf-8")
    req = urllib.request.Request(
        INGEST_ENDPOINT,
        data=body,
        method="POST",
        headers={"Content-Type": "application/json", "Accept": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return json.loads(resp.read().decode("utf-8"))


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--tabs", default=str(TABS_FILE), help=f"Path to tabs file (default {TABS_FILE})")
    ap.add_argument("--queue", action="store_true", help="Actually POST surviving URLs")
    ap.add_argument("--delay", type=float, default=0.15)
    ap.add_argument("--limit", type=int, default=None)
    ap.add_argument("--write-plan", default=None)
    args = ap.parse_args()

    tabs_path = Path(args.tabs)
    if not tabs_path.exists():
        print(f"tabs file not found: {tabs_path}", file=sys.stderr)
        return 2

    print(f"[1/5] reading {tabs_path}", file=sys.stderr)
    urls = parse_urls(tabs_path)
    print(f"      {len(urls)} url lines", file=sys.stderr)

    print("[2/5] loading nhentai DB-tags exclude set (artists/groups with posts ingested)", file=sys.stderr)
    nhentai_db_known = _load_from_postgres(PG_DSN)
    print(f"      {len(nhentai_db_known)} (kind,slug) pairs", file=sys.stderr)

    print("[3/5] loading creator rows by service", file=sys.stderr)
    creators_by_service = load_creator_slugs_by_service(PG_DSN)
    print(f"      {sum(len(v) for v in creators_by_service.values())} slugs across {len(creators_by_service)} services", file=sys.stderr)

    print("[4/5] loading in-flight tasks from server", file=sys.stderr)
    inflight = load_inflight_urls()
    print(f"      {len(inflight)} URLs currently queued or downloading", file=sys.stderr)

    print("[5/5] classifying", file=sys.stderr)
    keep: list[tuple[str, str]] = []
    drop_reasons: Counter[str] = Counter()
    skip_reasons: Counter[str] = Counter()
    handler_kept: Counter[str] = Counter()

    for url in urls:
        decision, reason = classify(url, nhentai_db_known, creators_by_service, inflight)
        if decision == "keep":
            keep.append((url, reason))
            handler_kept[reason] += 1
        elif decision == "drop":
            drop_reasons[reason] += 1
        else:
            skip_reasons[reason] += 1

    print(f"      kept (creators not yet finished): {len(keep)}", file=sys.stderr)
    print(f"      dropped (already in DB or in-flight):", file=sys.stderr)
    for r, n in drop_reasons.most_common():
        print(f"        {r:30s} {n}", file=sys.stderr)
    print(f"      skipped (non-creator URLs in file): {sum(skip_reasons.values())}", file=sys.stderr)
    print(f"      kept by handler:", file=sys.stderr)
    for h, n in handler_kept.most_common():
        print(f"        {h:25s} {n}", file=sys.stderr)

    if args.write_plan:
        Path(args.write_plan).write_text("\n".join(u for u, _ in keep) + "\n", encoding="utf-8")
        print(f"      wrote plan to {args.write_plan}", file=sys.stderr)

    if not args.queue:
        print("dry-run; pass --queue to POST these.", file=sys.stderr)
        for u, h in keep[:25]:
            print(f"{h:25s} {u}")
        if len(keep) > 25:
            print(f"... and {len(keep) - 25} more")
        return 0

    to_queue = keep[: args.limit] if args.limit else keep
    print(f"queueing {len(to_queue)} URLs (delay {args.delay}s)", file=sys.stderr)
    ok = fail = dup = 0
    fails: list[tuple[str, str]] = []
    for i, (url, handler) in enumerate(to_queue, 1):
        try:
            resp = post(url)
            if resp.get("duplicate"):
                dup += 1
                tag = "DUP"
            else:
                ok += 1
                tag = "OK "
            print(f"[{i:4d}/{len(to_queue)}] {tag} {handler:25s} {(resp.get('task_id','-') or '-')[:8]} {url}")
        except urllib.error.HTTPError as e:
            body = e.read().decode("utf-8", errors="replace")[:200]
            fail += 1
            fails.append((url, f"HTTP {e.code}: {body}"))
            print(f"[{i:4d}/{len(to_queue)}] ERR {handler:25s} HTTP {e.code} {url}", file=sys.stderr)
        except Exception as e:  # noqa: BLE001
            fail += 1
            fails.append((url, str(e)))
            print(f"[{i:4d}/{len(to_queue)}] ERR {handler:25s} {e} {url}", file=sys.stderr)
        if args.delay > 0 and i < len(to_queue):
            time.sleep(args.delay)

    print(f"\nqueued: ok={ok} dup={dup} fail={fail}", file=sys.stderr)
    if fails:
        print("first failures:", file=sys.stderr)
        for url, msg in fails[:10]:
            print(f"  {url}\n    {msg}", file=sys.stderr)
    return 0 if fail == 0 else 1


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