#!/usr/bin/env python3
"""
Discover TikTok videos (by hashtag, keyword search, or account) and their
comments via the TikHub TikTok APP v3 API, and save in the schema
run_jtbd_pipeline.py expects: post_url,author,comment,source_type,date,upvotes,downvotes.

Endpoints used (confirmed against TikHub's live OpenAPI spec at
https://api.tikhub.io/openapi.json -- not guessed):
  GET /api/v1/tiktok/app/v3/fetch_hashtag_search_result   (resolve hashtag name -> ch_id)
  GET /api/v1/tiktok/app/v3/fetch_hashtag_video_list       (ch_id, cursor)
  GET /api/v1/tiktok/app/v3/fetch_video_search_result      (keyword, offset)
  GET /api/v1/tiktok/app/v3/fetch_user_post_videos_v3      (unique_id, max_cursor)
  GET /api/v1/tiktok/app/v3/fetch_video_comments           (aweme_id, cursor)

fetch_hashtag_search_result's response shape (data.challenge_list[].challenge_info
with cha_name/cid) is already confirmed live in hashtag_stats_providers/tiktok.py.
The other four endpoints' response shapes are NOT published in the OpenAPI spec
(no example bodies) -- item/list field names below are best-effort guesses from
the same "aweme" API family, corrected against real responses once live (see
the design doc's Testing section).

Resumable: the per-topic video-listing CSV tracks a `scraped` flag per video,
and can be re-run with a different --mode against the same topic to union in
more videos (deduped by aweme_id) without re-fetching already-scraped ones.

Usage:
  python scrape_tiktok_via_tikhub.py declutter --mode hashtag --hashtag declutter --limit 100
  python scrape_tiktok_via_tikhub.py declutter --mode search --keyword "decluttering tips" --limit 100
  python scrape_tiktok_via_tikhub.py declutter --mode account --handles user1,user2 --limit 100
"""

import argparse
import csv
import json
import os
import time
from datetime import datetime, timezone
from pathlib import Path

from backend_client import REQUEST_DELAY_SECONDS, call_backend

VIDEO_LIST_CANDIDATE_PATHS = ["data.aweme_list", "data.search_item_list", "data.data", "data.videos", "data.itemList", "data.aweme_infos"]
COMMENT_LIST_CANDIDATE_PATHS = ["data.comments", "data.comment_list", "data.data"]

VIDEO_CSV_FIELDNAMES = ["aweme_id", "topic", "caption", "author", "like_count", "create_time", "scraped"]
COMMENT_CSV_FIELDNAMES = ["post_url", "author", "comment", "source_type", "date", "upvotes", "downvotes"]

MAX_DISCOVERY_PAGES = 25
MAX_PAGES_PER_HANDLE = 10
MAX_COMMENT_PAGES_PER_VIDEO = 15


def dig(data, dotted_path):
    """Walk a dotted key path through nested dicts. Returns None if any level is missing."""
    node = data
    for part in dotted_path.split("."):
        if not isinstance(node, dict) or part not in node:
            return None
        node = node[part]
    return node


def first_present(item, keys, default=""):
    """Return the first non-empty value among `keys` in dict `item`."""
    for key in keys:
        value = item.get(key)
        if value not in (None, ""):
            return value
    return default


def extract_list(raw_response, candidate_paths):
    """Try each candidate dotted path in order; return the first non-empty list found."""
    for path in candidate_paths:
        node = dig(raw_response, path)
        if isinstance(node, list) and node:
            return node
    return []


def pick_hashtag_id(raw_response, hashtag_name):
    """Find the ch_id for a hashtag name in a fetch_hashtag_search_result response.
    Response shape (data.challenge_list[].challenge_info.{cha_name,cid,...}) is
    already confirmed live by hashtag_stats_providers/tiktok.py -- not a guess.
    Prefers an exact case-insensitive name match; falls back to the first result.
    Returns None if the response has no hashtags at all."""
    challenges = dig(raw_response, "data.challenge_list") or []
    target = hashtag_name.strip().lower().lstrip("#")

    for challenge in challenges:
        info = challenge.get("challenge_info", {})
        if info.get("cha_name", "").lower() == target:
            return info.get("cid")

    if challenges:
        return challenges[0].get("challenge_info", {}).get("cid")
    return None


def format_create_time(value):
    """TikTok create_time is expected to be a Unix epoch (seconds). Returns ISO
    8601 UTC, or '' if the value is missing/unparseable."""
    try:
        return datetime.fromtimestamp(int(value), tz=timezone.utc).isoformat()
    except (TypeError, ValueError):
        return ""


def parse_video_item(item):
    """Extract the fields we need from one raw video-list item. Field names are
    best-guess -- corrected against real live responses in Task 4.

    Confirmed live: fetch_hashtag_video_list and fetch_user_post_videos_v3 both
    return flat aweme items directly; fetch_video_search_result instead wraps
    the real aweme fields one level deeper under an "aweme_info" key -- unwrap
    it here so both shapes produce the same result."""
    if isinstance(item.get("aweme_info"), dict):
        item = item["aweme_info"]

    author_obj = item.get("author") if isinstance(item.get("author"), dict) else {}
    stats_obj = item.get("statistics") if isinstance(item.get("statistics"), dict) else {}
    return {
        "aweme_id": str(first_present(item, ["aweme_id", "id", "video_id"])),
        "caption": first_present(item, ["desc", "caption", "title"]),
        "author": first_present(author_obj, ["unique_id", "nickname"]) or first_present(item, ["author_name"]),
        "like_count": first_present(stats_obj, ["digg_count", "like_count"]) or first_present(item, ["digg_count", "like_count"]),
        "create_time": first_present(item, ["create_time", "createTime"]),
    }


def parse_comment_item(item):
    """Extract the fields we need from one raw comment-list item. Field names are
    best-guess -- corrected against real live responses in Task 4."""
    user_obj = item.get("user") if isinstance(item.get("user"), dict) else {}
    return {
        "text": first_present(item, ["text", "content", "comment"]),
        "author": first_present(user_obj, ["unique_id", "nickname"]) or first_present(item, ["author_name"]),
        "like_count": first_present(item, ["digg_count", "like_count"]),
        "create_time": first_present(item, ["create_time", "createTime"]),
    }


def fetch_hashtag_videos(hashtag, region, limit):
    search_resp = call_backend("/api/v1/tiktok/app/v3/fetch_hashtag_search_result", {"keyword": hashtag})
    if search_resp is None:
        raise SystemExit(f"❌ Could not resolve hashtag '#{hashtag}' -- fetch_hashtag_search_result failed after retries")
    ch_id = pick_hashtag_id(search_resp, hashtag)
    if not ch_id:
        raise SystemExit(f"❌ No hashtag matching '#{hashtag}' found -- not writing any output")

    videos = []
    cursor = 0
    for _ in range(MAX_DISCOVERY_PAGES):
        if len(videos) >= limit:
            break
        params = {"ch_id": ch_id, "cursor": cursor, "count": 20, "region": region}
        raw = call_backend("/api/v1/tiktok/app/v3/fetch_hashtag_video_list", params)
        if raw is None:
            break
        items = extract_list(raw, VIDEO_LIST_CANDIDATE_PATHS)
        if not items:
            break
        videos.extend(parse_video_item(i) for i in items)
        print(f"  📄 Collected {len(videos)} videos so far...")
        if not dig(raw, "data.has_more"):
            break
        cursor = dig(raw, "data.cursor") or (cursor + len(items))
        time.sleep(REQUEST_DELAY_SECONDS)

    return videos[:limit]


def fetch_search_videos(keyword, region, limit):
    videos = []
    offset = 0
    for _ in range(MAX_DISCOVERY_PAGES):
        if len(videos) >= limit:
            break
        params = {"keyword": keyword, "offset": offset, "count": 20, "region": region}
        raw = call_backend("/api/v1/tiktok/app/v3/fetch_video_search_result", params)
        if raw is None:
            break
        items = extract_list(raw, VIDEO_LIST_CANDIDATE_PATHS)
        if not items:
            break
        videos.extend(parse_video_item(i) for i in items)
        print(f"  📄 Collected {len(videos)} videos so far...")
        if not dig(raw, "data.has_more"):
            break
        offset += len(items)
        time.sleep(REQUEST_DELAY_SECONDS)

    return videos[:limit]


def fetch_account_videos(handles, limit):
    videos = []
    per_handle_limit = max(limit // max(len(handles), 1), 1)

    for handle in handles:
        collected = 0
        max_cursor = 0
        for _ in range(MAX_PAGES_PER_HANDLE):
            if collected >= per_handle_limit:
                break
            params = {"unique_id": handle, "max_cursor": max_cursor, "count": 20}
            raw = call_backend("/api/v1/tiktok/app/v3/fetch_user_post_videos_v3", params)
            if raw is None:
                break
            items = extract_list(raw, VIDEO_LIST_CANDIDATE_PATHS)
            if not items:
                break
            new_videos = [parse_video_item(i) for i in items]
            videos.extend(new_videos)
            collected += len(new_videos)
            print(f"  📄 @{handle}: collected {collected} videos so far...")
            if not dig(raw, "data.has_more"):
                break
            max_cursor = dig(raw, "data.max_cursor") or (max_cursor + len(items))
            time.sleep(REQUEST_DELAY_SECONDS)

    return videos[:limit]


def load_existing_videos(videos_csv):
    if not os.path.exists(videos_csv):
        return {}
    with open(videos_csv, "r", encoding="utf-8", newline="") as f:
        return {row["aweme_id"]: row for row in csv.DictReader(f)}


def merge_videos(existing, discovered, topic):
    """Add newly-discovered videos not already present (by aweme_id). Existing
    entries (and their `scraped` flag) are left untouched."""
    for video in discovered:
        if video["aweme_id"] in existing:
            continue
        existing[video["aweme_id"]] = {**video, "topic": topic, "scraped": ""}
    return existing


def save_videos(videos_by_id, videos_csv):
    with open(videos_csv, "w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=VIDEO_CSV_FIELDNAMES)
        writer.writeheader()
        writer.writerows(videos_by_id.values())


def video_url(aweme_id):
    return f"https://www.tiktok.com/@_/video/{aweme_id}"


def caption_row(video):
    return {
        "post_url": video_url(video["aweme_id"]),
        "author": video["author"],
        "comment": video["caption"],
        "source_type": "caption",
        "date": format_create_time(video["create_time"]),
        "upvotes": video["like_count"],
        "downvotes": "",
    }


def comment_row(video, parsed_comment):
    return {
        "post_url": video_url(video["aweme_id"]),
        "author": parsed_comment["author"],
        "comment": parsed_comment["text"],
        "source_type": "comment",
        "date": format_create_time(parsed_comment["create_time"]),
        "upvotes": parsed_comment["like_count"],
        "downvotes": "",
    }


def fetch_comments_for_video(aweme_id, debug_dump_path=None):
    rows = []
    cursor = 0
    for page in range(MAX_COMMENT_PAGES_PER_VIDEO):
        raw = call_backend("/api/v1/tiktok/app/v3/fetch_video_comments", {"aweme_id": aweme_id, "cursor": cursor, "count": 20})
        if raw is None:
            break
        if debug_dump_path and page == 0:
            Path(debug_dump_path).write_text(json.dumps(raw, indent=2, ensure_ascii=False))
        items = extract_list(raw, COMMENT_LIST_CANDIDATE_PATHS)
        if not items:
            break
        rows.extend(parse_comment_item(i) for i in items)
        if not dig(raw, "data.has_more"):
            break
        cursor = dig(raw, "data.cursor") or (cursor + len(items))
        time.sleep(REQUEST_DELAY_SECONDS)

    return rows


def append_comments_csv(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_CSV_FIELDNAMES)
        if write_header:
            writer.writeheader()
        writer.writerows(rows)


def build_arg_parser():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("topic", help="Topic name -- groups videos from multiple discovery runs/modes into one output")
    parser.add_argument("--mode", required=True, choices=["hashtag", "search", "account"])
    parser.add_argument("--hashtag", help="Hashtag name (without '#'), required for --mode hashtag")
    parser.add_argument("--keyword", help="Search keyword, required for --mode search")
    parser.add_argument("--handles", help="Comma-separated TikTok usernames, required for --mode account")
    parser.add_argument("--limit", type=int, default=100, help="Number of videos to discover this run (default: 100)")
    parser.add_argument("--region", default="US", help="Region code, used by hashtag/search modes (default: US)")
    return parser


def main():
    args = build_arg_parser().parse_args()

    if args.mode == "hashtag" and not args.hashtag:
        raise SystemExit("❌ --hashtag is required for --mode hashtag")
    if args.mode == "search" and not args.keyword:
        raise SystemExit("❌ --keyword is required for --mode search")
    if args.mode == "account" and not args.handles:
        raise SystemExit("❌ --handles is required for --mode account")

    videos_csv = f"tt_{args.topic}_tikhub_videos.csv"
    comments_csv = f"tt_{args.topic}_tikhub_comments.csv"

    print(f"🔍 Discovering videos for topic '{args.topic}' via --mode {args.mode}...")
    if args.mode == "hashtag":
        discovered = fetch_hashtag_videos(args.hashtag, args.region, args.limit)
    elif args.mode == "search":
        discovered = fetch_search_videos(args.keyword, args.region, args.limit)
    else:
        handles = [h.strip() for h in args.handles.split(",") if h.strip()]
        discovered = fetch_account_videos(handles, args.limit)

    if not discovered:
        raise SystemExit(
            f"❌ Got 0 videos for --mode {args.mode} -- treating this as a failed fetch, not empty "
            f"results. Not touching {videos_csv}."
        )

    existing = load_existing_videos(videos_csv)
    before = len(existing)
    videos_by_id = merge_videos(existing, discovered, args.topic)
    save_videos(videos_by_id, videos_csv)
    print(f"📄 {len(videos_by_id) - before} new video(s), {len(videos_by_id)} total in {videos_csv}")

    remaining = [v for v in videos_by_id.values() if v.get("scraped") != "True"]
    print(f"📋 {len(remaining)}/{len(videos_by_id)} videos left to scrape comments for\n")

    dumped_debug_sample = False
    for idx, video in enumerate(remaining, 1):
        print(f"[{idx}/{len(remaining)}] Fetching comments: {video['aweme_id']} ({video['caption'][:60]})")

        append_comments_csv([caption_row(video)], comments_csv)

        debug_path = None
        if not dumped_debug_sample:
            debug_path = f"tt_{args.topic}_{video['aweme_id']}_comments_raw.json"
            dumped_debug_sample = True

        parsed_comments = fetch_comments_for_video(video["aweme_id"], debug_dump_path=debug_path)
        append_comments_csv([comment_row(video, c) for c in parsed_comments], comments_csv)
        print(
            f"  💾 Saved 1 caption + {len(parsed_comments)} comment(s) -> {comments_csv}"
            + (f" (raw dump: {debug_path})" if debug_path else "")
        )

        video["scraped"] = "True"
        save_videos(videos_by_id, videos_csv)
        time.sleep(REQUEST_DELAY_SECONDS)

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


if __name__ == "__main__":
    main()