corpus.py

#!/usr/bin/env python3
"""pakka_corpus CLI — pull, store, surface Reddit threads for taxonomy refinement."""
from __future__ import annotations

import argparse
import json
import logging
import re
import sys
import time
from pathlib import Path
from typing import Optional

import yaml

from analyze import format_surface, keyword_hits
from arctic_shift import ArcticShiftClient, ArcticShiftError
from pullpush import PullPushClient, PullPushError
from reddit_api import RedditClient, RedditError, parse_thread_url
from store import Store, _extract_media_json

ROOT = Path(__file__).parent
CONFIG_PATH = ROOT / "config.yaml"
SEEDS_PATH = ROOT / "seeds.txt"
DB_PATH = ROOT / "corpus.db"

log = logging.getLogger("corpus")


def load_config() -> dict:
    if not CONFIG_PATH.exists():
        sys.exit(f"config.yaml not found at {CONFIG_PATH}")
    with open(CONFIG_PATH, "r") as f:
        return yaml.safe_load(f) or {}


_SINCE_RE = re.compile(r"^(\d+)([dhmw])$", re.IGNORECASE)


def parse_since(since: Optional[str]) -> Optional[int]:
    if not since or since.lower() == "all":
        return None
    m = _SINCE_RE.match(since)
    if not m:
        raise ValueError(f"bad --since value: {since} (expected like 30d, 12h, 2w, 'all')")
    n, unit = int(m.group(1)), m.group(2).lower()
    seconds = {"h": 3600, "d": 86400, "w": 604800, "m": 2592000}[unit]
    return int(time.time()) - n * seconds


def build_clients(cfg: dict) -> tuple[ArcticShiftClient, PullPushClient, RedditClient]:
    rl = cfg.get("rate_limit", {}) or {}
    arctic = ArcticShiftClient(
        delay_seconds=float(rl.get("arctic_shift_delay_seconds", 1.0)),
        max_retries=int(rl.get("max_retries", 3)),
        backoff_base=int(rl.get("backoff_base", 2)),
    )
    pp = PullPushClient(
        delay_seconds=float(rl.get("pullpush_delay_seconds", 1.2)),
        max_retries=int(rl.get("max_retries", 3)),
        backoff_base=int(rl.get("backoff_base", 2)),
    )
    rd = RedditClient(delay_seconds=float(rl.get("reddit_delay_seconds", 2.0)))
    return arctic, pp, rd


def get_subreddits_to_pull(cfg: dict, store: Store) -> list[str]:
    names: list[str] = []
    seen: set[str] = set()
    for entry in cfg.get("subreddits", []) or []:
        n = entry.get("name") if isinstance(entry, dict) else str(entry)
        if n and n not in seen:
            names.append(n); seen.add(n)
    for n in store.selected_subreddits():
        if n not in seen:
            names.append(n); seen.add(n)
    return names


def cmd_init(args):
    Store(DB_PATH).close()
    print(f"Initialized SQLite DB at {DB_PATH}")
    print(f"Config:  {CONFIG_PATH}")
    print(f"Seeds:   {SEEDS_PATH}")


def cmd_discover(args):
    cfg = load_config()
    _, _pp, reddit = build_clients(cfg)
    store = Store(DB_PATH)
    try:
        results = reddit.search_subreddits(args.query, limit=args.limit)
    except RedditError as e:
        sys.exit(f"Reddit search failed: {e}")
    if not results:
        print("No subreddits matched."); return
    print(f"\nFound {len(results)} subreddits for query: {args.query!r}\n")
    print(f"{'subscribers':>11}  {'NSFW':>5}  name  —  description")
    print("-" * 80)
    for r in results:
        store.upsert_subreddit(r, discovered_via=f"search:{args.query}")
        nsfw = "yes" if r["over_18"] else ""
        desc = (r["description"][:80] + "…") if len(r["description"]) > 80 else r["description"]
        print(f"{r['subscribers']:>11,}  {nsfw:>5}  {r['name']}  —  {desc}")
    store.commit(); store.close()
    print(f"\nTo include in pulls: python corpus.py select <name>")


def cmd_select(args):
    s = Store(DB_PATH); s.set_selected(args.name, True); s.close()
    print(f"Marked r/{args.name} as selected.")


def cmd_unselect(args):
    s = Store(DB_PATH); s.set_selected(args.name, False); s.close()
    print(f"Unmarked r/{args.name}.")


def cmd_list_subs(args):
    cfg = load_config()
    store = Store(DB_PATH)
    in_config = [(e.get("name") if isinstance(e, dict) else str(e))
                 for e in (cfg.get("subreddits") or [])]
    rows = store.list_subreddits()
    store.close()
    print("\nFrom config.yaml:")
    for n in in_config:
        print(f"  - {n}")
    if rows:
        print("\nDiscovered subreddits (* = selected):")
        for r in rows:
            mark = "*" if r["selected"] else " "
            subs = r["subscribers"] or 0
            desc = (r["description"] or "")[:60]
            print(f"  {mark} r/{r['name']:30s} {subs:>10,} subs   {desc}")
    else:
        print("\n(no discovered subreddits yet — run `discover --query ...`)")


def _pull_submissions_with_fallback(arctic, pp, sub, since_ts, max_posts):
    """Yield (item, source). Arctic Shift primary; PullPush fallback on first-call failure."""
    yielded = 0
    try:
        for item in arctic.paginate_submissions(sub, since_ts=since_ts, max_items=max_posts):
            yield item, "arctic_shift"
            yielded += 1
        return
    except ArcticShiftError as e:
        if yielded > 0:
            log.warning("Arctic Shift failed mid-stream (%d posts): %s", yielded, e)
            return
        log.warning("Arctic Shift failed, falling back to PullPush: %s", e)
    for item in pp.paginate_submissions(sub, since_ts=since_ts, max_items=max_posts):
        yield item, "pullpush"


def _pull_comments_with_fallback(arctic, pp, sub, since_ts, max_comments):
    """Yield (item, source). Arctic Shift primary; PullPush fallback on first-call failure."""
    yielded = 0
    try:
        for item in arctic.paginate_comments_for_subreddit(sub, since_ts=since_ts, max_items=max_comments):
            yield item, "arctic_shift"
            yielded += 1
        return
    except ArcticShiftError as e:
        if yielded > 0:
            log.warning("Arctic Shift failed mid-stream (%d comments): %s", yielded, e)
            return
        log.warning("Arctic Shift failed, falling back to PullPush: %s", e)
    for item in pp.paginate_comments_for_subreddit(sub, since_ts=since_ts, max_items=max_comments):
        yield item, "pullpush"


def cmd_pull(args):
    cfg = load_config()
    arctic, pp, _ = build_clients(cfg)
    store = Store(DB_PATH)
    window = cfg.get("default_window", {}) or {}
    since = args.since or window.get("since", "365d")
    since_ts = parse_since(since)
    max_posts = args.max_posts or int(window.get("max_posts_per_sub", 2000))
    max_comments = args.max_comments or int(window.get("max_comments_per_sub", 10000))
    subs = [args.subreddit] if args.subreddit else get_subreddits_to_pull(cfg, store)
    if not subs:
        sys.exit("No subreddits configured. Add to config.yaml or run `discover` + `select`.")
    print(f"Pulling from {len(subs)} subreddit(s); since={since}, max_posts={max_posts}, max_comments={max_comments}")
    print(f"Source: Arctic Shift (primary) → PullPush (fallback)")
    for sub in subs:
        print(f"\n=== r/{sub} ===")
        run_id = store.start_run("subreddit_posts", sub, {"since": since, "max": max_posts})
        added = 0
        try:
            for i, (item, src) in enumerate(_pull_submissions_with_fallback(arctic, pp, sub, since_ts, max_posts), 1):
                if store.upsert_post(item, source=src):
                    added += 1
                if i % 200 == 0:
                    store.commit(); print(f"  posts: {i} fetched, {added} stored")
            store.commit()
            store.finish_run(run_id, posts_added=added, comments_added=0)
            print(f"  posts: done. {added} stored")
        except Exception as e:
            store.finish_run(run_id, posts_added=added, comments_added=0, error=str(e))
            print(f"  posts: FAILED — {e}"); continue
        run_id = store.start_run("subreddit_comments", sub, {"since": since, "max": max_comments})
        added = 0
        try:
            for i, (item, src) in enumerate(_pull_comments_with_fallback(arctic, pp, sub, since_ts, max_comments), 1):
                if store.upsert_comment(item, source=src):
                    added += 1
                if i % 500 == 0:
                    store.commit(); print(f"  comments: {i} fetched, {added} stored")
            store.commit()
            store.finish_run(run_id, posts_added=0, comments_added=added)
            print(f"  comments: done. {added} stored")
        except Exception as e:
            store.finish_run(run_id, posts_added=0, comments_added=added, error=str(e))
            print(f"  comments: FAILED — {e}")
    store.close()
    print("\nPull complete.")


def cmd_pull_thread(args):
    cfg = load_config()
    arctic, pp, reddit = build_clients(cfg)
    store = Store(DB_PATH)
    parsed = parse_thread_url(args.url)
    if not parsed:
        sys.exit(f"Could not parse thread URL: {args.url}")
    sub, post_id = parsed
    print(f"Thread: r/{sub} post_id={post_id}")
    run_id = store.start_run("thread", post_id, {"url": args.url})
    posts_added = 0; comments_added = 0; err = None
    try:
        thread = reddit.get_thread(post_id)
        if thread and isinstance(thread, list) and len(thread) >= 1:
            post_listing = thread[0].get("data", {}).get("children", [])
            if post_listing:
                post_data = post_listing[0].get("data", {})
                if store.upsert_post(post_data, source="reddit"):
                    posts_added += 1
    except Exception as e:
        log.warning("live thread fetch failed; continuing with archive only: %s", e)
    # Comments: Arctic Shift primary, PullPush fallback
    try:
        fetched = 0
        for client, src in [(arctic, "arctic_shift"), (pp, "pullpush")]:
            try:
                for i, item in enumerate(client.paginate_comments_for_post(post_id, max_items=20000), 1):
                    if store.upsert_comment(item, source=src):
                        comments_added += 1
                    if i % 500 == 0:
                        store.commit(); print(f"  comments: {i} fetched, {comments_added} stored")
                    fetched += 1
                store.commit()
                break  # success — don't try fallback
            except (ArcticShiftError, PullPushError) as e:
                if fetched > 0:
                    err = str(e); break  # failed mid-stream
                log.warning("%s failed for comments, trying next source: %s", src, e)
    except Exception as e:
        err = str(e)
    store.finish_run(run_id, posts_added=posts_added, comments_added=comments_added, error=err)
    store.close()
    print(f"Done. posts_added={posts_added} comments_added={comments_added}" + (f" error={err}" if err else ""))


def cmd_pull_seeds(args):
    if not SEEDS_PATH.exists():
        sys.exit(f"seeds.txt not found at {SEEDS_PATH}")
    urls = []
    with open(SEEDS_PATH, "r") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"): continue
            urls.append(line)
    if not urls:
        print("No seed URLs in seeds.txt (lines starting with # are comments)."); return
    print(f"Pulling {len(urls)} seed thread(s)...")
    for url in urls:
        print(f"\n--- {url}")
        try:
            cmd_pull_thread(argparse.Namespace(url=url))
        except SystemExit as e:
            print(f"  skipped: {e}")


def cmd_stats(args):
    s = Store(DB_PATH); st = s.stats(); s.close()
    print(f"\nposts:    {st['posts']:,}")
    print(f"comments: {st['comments']:,}")
    print(f"runs:     {st['total_runs']:,}")
    if st["posts_by_subreddit"]:
        print("\nPosts by subreddit:")
        for r in st["posts_by_subreddit"]:
            print(f"  {r['n']:>8,}  r/{r['subreddit']}")
    if st["comments_by_subreddit"]:
        print("\nComments by subreddit:")
        for r in st["comments_by_subreddit"]:
            print(f"  {r['n']:>8,}  r/{r['subreddit']}")
    if st["recent_runs"]:
        print("\nRecent runs:")
        for r in st["recent_runs"]:
            status = "ok" if not r["error"] else f"ERR: {r['error'][:40]}"
            print(f"  {r['kind']:22s} {r['target']:25s} +{r['posts_added']}p +{r['comments_added']}c  {status}")


def cmd_scan(args):
    s = Store(DB_PATH)
    rows = s.search_comments_like(args.pattern, subreddit=args.subreddit, limit=args.limit)
    s.close()
    if not rows:
        print(f"No hits for {args.pattern!r}."); return
    print(f"\n{len(rows)} hit(s) for {args.pattern!r}:\n")
    for r in rows:
        body = (r["body"] or "").replace("\n", " ")[:140]
        when = time.strftime("%Y-%m-%d", time.gmtime(r["created_utc"]))
        print(f"  [{when}] r/{r['subreddit']}  u/{r['author']}  (score={r['score']})")
        print(f"     {body}")
        if r["permalink"]:
            print(f"     https://www.reddit.com{r['permalink']}")
        print()


def cmd_surface(args):
    cfg = load_config()
    tax = cfg.get("taxonomy_keywords", {}) or {}
    if not tax:
        sys.exit("No taxonomy_keywords in config.yaml")
    store = Store(DB_PATH)
    print("Scanning corpus for taxonomy keyword hits...")
    results = keyword_hits(store, tax, subreddit=args.subreddit)
    print(format_surface(results, store))
    store.close()


def cmd_backfill_media(args):
    """Re-fetch non-self posts from Reddit to fill media_json (gallery images, videos)."""
    cfg = load_config()
    _, _, reddit = build_clients(cfg)
    store = Store(DB_PATH)
    rows = store.conn.execute(
        "SELECT id FROM posts WHERE media_json IS NULL AND is_self = 0"
    ).fetchall()
    ids = [r[0] for r in rows]
    if not ids:
        print("All non-self posts already have media_json set. Nothing to do.")
        store.close()
        return
    print(f"Backfilling media for {len(ids)} posts via Reddit /api/info…")
    updated = 0
    batch_size = 100
    for i in range(0, len(ids), batch_size):
        batch = ids[i:i + batch_size]
        try:
            items = reddit.fetch_posts_by_ids(batch)
            for item in items:
                pid = item.get("id") or ""
                if not pid:
                    continue
                media_json = _extract_media_json(item)
                if media_json:
                    store.conn.execute(
                        "UPDATE posts SET media_json = ? WHERE id = ? AND media_json IS NULL",
                        (media_json, pid),
                    )
                    updated += 1
            store.commit()
            done = min(i + batch_size, len(ids))
            print(f"  {done}/{len(ids)} checked, {updated} updated with media")
        except RedditError as e:
            print(f"  Reddit error on batch {i // batch_size + 1}: {e}")
            time.sleep(5)
    store.close()
    print(f"\nDone. {updated} posts updated with media.")


def cmd_export(args):
    s = Store(DB_PATH); n = 0
    out = Path(args.output)
    with open(out, "w") as f:
        kinds = (["posts", "comments"] if args.kind == "both" else [args.kind])
        for kind in kinds:
            for row in s.iter_all(kind, subreddit=args.subreddit):
                row["_kind"] = kind
                f.write(json.dumps(row, ensure_ascii=False) + "\n")
                n += 1
    s.close()
    print(f"Wrote {n} rows → {out}")


def build_parser():
    p = argparse.ArgumentParser(prog="corpus.py", description=__doc__)
    p.add_argument("--verbose", "-v", action="store_true")
    sub = p.add_subparsers(dest="cmd", required=True)
    sub.add_parser("init").set_defaults(func=cmd_init)
    d = sub.add_parser("discover"); d.add_argument("--query", "-q", required=True); d.add_argument("--limit", "-n", type=int, default=25); d.set_defaults(func=cmd_discover)
    s = sub.add_parser("select"); s.add_argument("name"); s.set_defaults(func=cmd_select)
    u = sub.add_parser("unselect"); u.add_argument("name"); u.set_defaults(func=cmd_unselect)
    sub.add_parser("list-subs").set_defaults(func=cmd_list_subs)
    pl = sub.add_parser("pull"); pl.add_argument("--subreddit", "-s"); pl.add_argument("--since"); pl.add_argument("--max-posts", type=int); pl.add_argument("--max-comments", type=int); pl.set_defaults(func=cmd_pull)
    pt = sub.add_parser("pull-thread"); pt.add_argument("url"); pt.set_defaults(func=cmd_pull_thread)
    sub.add_parser("pull-seeds").set_defaults(func=cmd_pull_seeds)
    sub.add_parser("stats").set_defaults(func=cmd_stats)
    sc = sub.add_parser("scan"); sc.add_argument("pattern"); sc.add_argument("--subreddit", "-s"); sc.add_argument("--limit", "-n", type=int, default=50); sc.set_defaults(func=cmd_scan)
    surf = sub.add_parser("surface"); surf.add_argument("--subreddit", "-s"); surf.set_defaults(func=cmd_surface)
    ex = sub.add_parser("export"); ex.add_argument("--kind", choices=["posts","comments","both"], default="comments"); ex.add_argument("--subreddit", "-s"); ex.add_argument("--output", "-o", required=True); ex.set_defaults(func=cmd_export)
    sub.add_parser("backfill-media").set_defaults(func=cmd_backfill_media)
    return p


def main():
    args = build_parser().parse_args()
    level = logging.DEBUG if args.verbose else logging.INFO
    logging.basicConfig(level=level, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
    args.func(args)


if __name__ == "__main__":
    main()