#!/usr/bin/env python3
"""
Scrape a subreddit's top posts + comments via the TikHub Reddit APP API
(https://api.tikhub.io) instead of hitting reddit.com directly, and save in
the schema run_jtbd_pipeline.py expects: post_url,author,comment,date,upvotes,downvotes.

Endpoints used (confirmed against TikHub's live OpenAPI spec at
https://api.tikhub.io/openapi.json -- not guessed):
  GET /api/v1/reddit/app/fetch_subreddit_feed  (subreddit_name, sort, after)
  GET /api/v1/reddit/app/fetch_dynamic_search  (query, search_type, sort, time_range, after)
  GET /api/v1/reddit/app/fetch_post_comments   (post_id, sort_type, after)
  GET /api/v1/reddit/app/fetch_comment_replies (post_id, cursor, sort_type)

With --query, posts come from fetch_dynamic_search using Reddit's native
"subreddit:<name> <query>" search syntax (confirmed live: results are scoped
to that subreddit) instead of fetch_subreddit_feed. Its response shape
differs from the feed's: each edge under
data.search.dynamic.components.main.edges wraps several posts in a
"children" list (not one post per edge), and pagination uses that same
component's pageInfo.hasNextPage/endCursor.

The feed response is a GraphQL-style edges/node/cells structure (Reddit's
internal app API shape, passed through by TikHub). Each post's "cells" list
carries typed cells identified by their key set: a MetadataCell
(createdAt, authorName), a TitleCell or TitleAndThumbnailCell (title), and
an ActionCell (score, commentCount). This was confirmed against a real
successful response; the comment-tree shape below is a best-effort parse
based on the same edges/node convention used everywhere else in this API --
the raw JSON for the first post's comments is always saved to
<post_id>_comments_raw.json for easy debugging if TikHub's actual shape
differs once you can test this live.

Resumable: the posts-listing CSV tracks a `scraped` flag per post.

Usage:
  python scrape_subreddit_via_tikhub.py quittingsmoking --limit 100 --sort TOP
  python scrape_subreddit_via_tikhub.py houseplants --query water --limit 500 --max-comment-pages 2
"""

import argparse
import csv
import os
import re
import time
from pathlib import Path

import requests
from dotenv import load_dotenv

load_dotenv()

BASE_URL = "https://api.tikhub.io"
API_KEY = os.environ["TIKHUB_API_KEY"]
HEADERS = {"Authorization": f"Bearer {API_KEY}"}

MAX_RETRIES = 4
RETRY_BACKOFF_SECONDS = 5.0
REQUEST_DELAY_SECONDS = 0.5
HTTP_OK = 200
HTTP_RATE_LIMITED = 429

POST_FIELDNAMES = ["post_url", "post_id", "title", "author", "score", "num_comments", "created_at", "scraped"]
COMMENT_FIELDNAMES = ["post_url", "author", "comment", "date", "upvotes", "downvotes"]


def call_tikhub(path, params):
    """GET a TikHub endpoint with retries. Returns parsed JSON or None on repeated failure."""
    url = f"{BASE_URL}{path}"
    delay = RETRY_BACKOFF_SECONDS
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            r = requests.get(url, headers=HEADERS, params=params, timeout=40)
        except requests.RequestException as e:
            print(f"  ⚠️ Network error on {path} (attempt {attempt}/{MAX_RETRIES}): {e}")
            time.sleep(delay)
            delay *= 2
            continue

        if r.status_code == HTTP_OK:
            return r.json()

        print(f"  ⚠️ {path} -> HTTP {r.status_code} (attempt {attempt}/{MAX_RETRIES}): {r.text[:200]}")
        if r.status_code == HTTP_RATE_LIMITED or 500 <= r.status_code < 600 or r.status_code == 400:
            time.sleep(delay)
            delay *= 2
            continue
        return None

    return None


def parse_feed_edges(feed_json):
    """Extract post entries from a fetch_subreddit_feed response.

    The response nests differently depending on request params observed live:
    with need_format it's data.subredditfeed.subredditV3, without it's
    data.subredditV3 directly -- handle both.
    """
    data = feed_json["data"]
    subreddit_v3 = data["subredditfeed"]["subredditV3"] if "subredditfeed" in data else data["subredditV3"]
    edges = subreddit_v3["elements"]["edges"]
    page_info = subreddit_v3["elements"].get("pageInfo", {})
    end_cursor = page_info.get("endCursor")

    posts = []
    for edge in edges:
        node = edge["node"]
        group_id = node.get("groupId", "")
        if not group_id.startswith("t3_"):
            continue  # skip non-post cells like "sortcell"

        title = ""
        author = ""
        created_at = ""
        score = ""
        num_comments = ""
        for cell in node.get("cells", []):
            if "authorName" in cell and "createdAt" in cell:
                author = cell.get("authorName", "").lstrip("u/")
                created_at = cell.get("createdAt", "")
            if "title" in cell and "isVisited" in cell:
                title = cell.get("title", "")
            elif "titleCell" in cell:
                title = cell.get("titleCell", {}).get("title", "")
            if "score" in cell and "commentCount" in cell:
                score = cell.get("score", "")
                num_comments = cell.get("commentCount", "")

        post_id36 = group_id.split("_", 1)[1]
        posts.append(
            {
                "post_url": f"https://www.reddit.com/comments/{post_id36}/",
                "post_id": group_id,
                "title": title,
                "author": author,
                "score": score,
                "num_comments": num_comments,
                "created_at": created_at,
                "scraped": "",
            }
        )
    return posts, end_cursor


def parse_search_edges(search_json):
    """Extract post entries from a fetch_dynamic_search (search_type=post) response.

    Shape differs from the feed: data.search.dynamic.components.main.edges is a
    list where each edge's node.children is itself a list of {"post": {...}}
    entries (several posts bundled per edge, not one post per edge). Pagination
    cursor lives on that same "main" component, not per-edge.
    """
    main = search_json["data"]["search"]["dynamic"]["components"]["main"]
    end_cursor = main.get("pageInfo", {}).get("endCursor") if main.get("pageInfo", {}).get("hasNextPage") else None

    posts = []
    for edge in main.get("edges", []):
        for child in edge.get("node", {}).get("children", []):
            post = child.get("post")
            if not post:
                continue
            post_id = post["id"]  # already "t3_..." here, unlike the feed's groupId
            post_id36 = post_id.split("_", 1)[1]
            posts.append(
                {
                    "post_url": f"https://www.reddit.com/comments/{post_id36}/",
                    "post_id": post_id,
                    "title": post.get("postTitle", ""),
                    "author": (post.get("authorInfo") or {}).get("name", ""),
                    "score": post.get("score", ""),
                    "num_comments": post.get("commentCount", ""),
                    "created_at": post.get("createdAt", ""),
                    "scraped": "",
                }
            )
    return posts, end_cursor


def load_or_fetch_posts(subreddit, limit, sort, posts_csv, query=None):
    if os.path.exists(posts_csv):
        with open(posts_csv, "r", encoding="utf-8", newline="") as f:
            posts = list(csv.DictReader(f))
        if not posts:
            raise SystemExit(f"❌ {posts_csv} exists but has 0 rows -- remove it and re-run.")
        print(f"🔄 Resuming from existing listing: {posts_csv} ({len(posts)} posts)")
        return posts

    if query:
        print(f"🔍 Searching {limit} posts (query={query!r}, sort={sort}) in r/{subreddit} via TikHub...")
    else:
        print(f"🔍 Fetching top {limit} posts (sort={sort}) from r/{subreddit} via TikHub...")
    posts = []
    after = None
    while len(posts) < limit:
        if query:
            params = {
                "query": f"subreddit:{subreddit} {query}",
                "search_type": "post",
                "sort": sort,
                "time_range": "all",
            }
            path = "/api/v1/reddit/app/fetch_dynamic_search"
        else:
            params = {"subreddit_name": subreddit, "sort": sort}
            path = "/api/v1/reddit/app/fetch_subreddit_feed"
        if after:
            params["after"] = after
        feed_json = call_tikhub(path, params)
        if feed_json is None:
            break
        page_posts, after = parse_search_edges(feed_json) if query else parse_feed_edges(feed_json)
        if not page_posts:
            break
        posts.extend(page_posts)
        print(f"  📄 Collected {len(posts)} posts so far...")
        if not after:
            break
        time.sleep(REQUEST_DELAY_SECONDS)

    if not posts:
        raise SystemExit(
            f"❌ Got 0 posts for r/{subreddit} from TikHub -- treating this as a failed "
            f"fetch, not an empty subreddit. Not writing {posts_csv} so a re-run will retry."
        )

    posts = posts[:limit]
    save_posts(posts, posts_csv)
    print(f"📄 Saved {len(posts)} posts to {posts_csv}")
    return posts


def save_posts(posts, posts_csv):
    with open(posts_csv, "w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=POST_FIELDNAMES)
        writer.writeheader()
        writer.writerows(posts)


# fetch_post_comments's own pageInfo never carries a cursor in practice (confirmed
# live: {"commentCount": N, "hasNextPage": true} with no cursor field). The real
# pagination mechanism is a "MoreComments"-style sentinel entry embedded *in the
# trees list itself*: a tree with node=None, a depth, and a `more` dict carrying
# a cursor. The depth-0 sentinel means "more top-level comments exist"; feeding
# its cursor to fetch_comment_replies returns another batch of real comments
# across several depths PLUS a fresh depth-0 sentinel to continue with. We keep
# walking that chain until it stops appearing or we hit the page cap below.
MAX_COMMENT_PAGES_PER_POST = 15


def parse_comment_tree_entry(tree, post_url):
    """Extract one row from a real (non-sentinel) commentForest.trees[] entry.
    Each entry is a flat depth-first node (depth/parentId reconstruct hierarchy);
    we don't need hierarchy for the flat CSV output."""
    node = tree.get("node")
    if not node:
        return None  # sentinel ("more") entry or a tombstoned/removed comment
    author = (node.get("authorInfo") or {}).get("name", "")
    body = (node.get("content") or {}).get("markdown", "")
    return {
        "post_url": post_url,
        "author": author,
        "comment": body,
        "date": node.get("createdAt", ""),
        "upvotes": node.get("score", ""),
        "downvotes": "",
    }


def find_top_level_more_cursor(trees):
    """Find the depth-0 'more comments' sentinel in a trees list, if any."""
    for tree in trees:
        if tree.get("node") is None and tree.get("depth") == 0 and tree.get("more"):
            return tree["more"].get("cursor")
    return None


def fetch_comments_for_post(post_id, post_url, debug_dump_path=None, max_pages=MAX_COMMENT_PAGES_PER_POST):
    rows = []

    first_page = call_tikhub("/api/v1/reddit/app/fetch_post_comments", {"post_id": post_id, "sort_type": "TOP"})
    if first_page is None:
        return rows

    if debug_dump_path:
        import json

        Path(debug_dump_path).write_text(json.dumps(first_page, indent=2))

    forest = (first_page.get("data") or {}).get("postInfoById") or {}
    trees = (forest.get("commentForest") or {}).get("trees", [])
    rows.extend(r for t in trees if (r := parse_comment_tree_entry(t, post_url)))
    cursor = find_top_level_more_cursor(trees)

    for _ in range(max_pages - 1):
        if not cursor:
            break
        time.sleep(REQUEST_DELAY_SECONDS)
        page_json = call_tikhub(
            "/api/v1/reddit/app/fetch_comment_replies",
            {"post_id": post_id, "cursor": cursor, "sort_type": "TOP"},
        )
        if page_json is None:
            break
        forest = (page_json.get("data") or {}).get("postInfoById") or {}
        trees = (forest.get("commentForest") or {}).get("trees", [])
        if not trees:
            break
        rows.extend(r for t in trees if (r := parse_comment_tree_entry(t, post_url)))
        cursor = find_top_level_more_cursor(trees)
    else:
        if cursor:
            print(f"  ⚠️ Hit max_pages={max_pages} cap for {post_id} -- more top-level comments remained")

    return rows


def append_comments(rows, comments_csv):
    if not rows:
        return
    write_header = not os.path.exists(comments_csv)
    with open(comments_csv, "a", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=COMMENT_FIELDNAMES)
        if write_header:
            writer.writeheader()
        writer.writerows(rows)


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("subreddit", help="Subreddit name, without 'r/'")
    ap.add_argument("--limit", type=int, default=100, help="Number of top posts to fetch (default: 100)")
    ap.add_argument("--sort", default=None,
                     help="Sort method. Feed mode: BEST, HOT, NEW, TOP, CONTROVERSIAL, RISING (default TOP). "
                          "Search mode (--query): RELEVANCE, HOT, TOP, NEW, COMMENTS (default NEW).")
    ap.add_argument("--query", default=None,
                     help="Search keyword to scope posts to (uses fetch_dynamic_search with "
                          "'subreddit:<name> <query>' instead of the plain subreddit feed)")
    ap.add_argument("--max-comment-pages", type=int, default=MAX_COMMENT_PAGES_PER_POST,
                     help=f"Comment pages per post: 1 fetch_post_comments call + N-1 fetch_comment_replies "
                          f"calls (default: {MAX_COMMENT_PAGES_PER_POST})")
    args = ap.parse_args()

    sort = args.sort or ("NEW" if args.query else "TOP")
    slug = f"{re.sub(r'[^a-z0-9]+', '_', args.query.lower()).strip('_')}_" if args.query else ""

    posts_csv = f"r_{args.subreddit}_tikhub_{slug}{sort.lower()}_posts.csv"
    comments_csv = f"r_{args.subreddit}_tikhub_{slug}{sort.lower()}_comments.csv"

    posts = load_or_fetch_posts(args.subreddit, args.limit, sort, posts_csv, query=args.query)

    remaining = [p for p in posts if p.get("scraped") != "True"]
    print(f"📋 {len(remaining)}/{len(posts)} posts left to scrape\n")

    dumped_debug_sample = False
    for idx, post in enumerate(posts, 1):
        if post.get("scraped") == "True":
            continue

        post_id = post["post_id"]
        url = post["post_url"]
        print(f"[{idx}/{len(posts)}] Fetching comments: {post_id} ({post.get('title', '')[:60]})")

        debug_path = None
        if not dumped_debug_sample:
            debug_path = f"{args.subreddit}_{post_id}_comments_raw.json"
            dumped_debug_sample = True

        rows = fetch_comments_for_post(post_id, url, debug_dump_path=debug_path, max_pages=args.max_comment_pages)
        append_comments(rows, comments_csv)
        print(f"  💾 Saved {len(rows)} comments -> {comments_csv}" + (f" (raw dump: {debug_path})" if debug_path else ""))

        post["scraped"] = "True"
        save_posts(posts, posts_csv)

        time.sleep(REQUEST_DELAY_SECONDS)

    print(f"\n✅ Done. Posts: {posts_csv} | Comments (ready for run_jtbd_pipeline.py): {comments_csv}")


if __name__ == "__main__":
    main()