"""从抖音链接 jsonl 里筛出 AI 相关的作品。

判定逻辑在 douyin/aiclass.py，分 platform / strong / topic / author 四档，
这里只负责流式读写、统计和落盘。

输出：
  *.txt    纯链接一行一个，直接喂给 pipeline.py --from-file
  *.jsonl  原记录 + ai_tier + ai_signals，字段全留着，方便复核和再切。
           这份体积接近源文件（实测 13967 条就有 169MB），而且随时能重跑
           生成，所以默认**不写**，要复核再加 --with-jsonl。

用法:
    python clean_ai.py -i 抖音链接20260827_40087.jsonl --stats
    python clean_ai.py -i 抖音链接20260827_40087.jsonl --tier platform,strong -o ai
    python clean_ai.py -i 抖音链接20260827_40087.jsonl --tier platform,strong -o ai --with-jsonl
    python clean_ai.py -i xxx.jsonl --tier platform,strong --exclude-forbidden -o ai
"""

from __future__ import annotations

# 移到 scripts/ tests/ 后，直接运行时脚本目录成了 sys.path[0]，
# 把项目根加回来，import douyin / pipeline 才找得到。
import sys as _sys
from pathlib import Path as _Path
_sys.path.insert(0, str(_Path(__file__).resolve().parent.parent))

import argparse
import json
import sys
from collections import Counter
from pathlib import Path

from douyin.aiclass import TIERS, classify


def iter_jsonl(path: Path):
    with open(path, "r", encoding="utf-8", errors="replace") as f:
        for n, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                yield n, json.loads(line)
            except ValueError:
                print(f"[跳过] 第 {n} 行不是合法 JSON", file=sys.stderr)


def main() -> int:
    ap = argparse.ArgumentParser(description="筛出 AI 相关作品")
    ap.add_argument("-i", "--input", required=True)
    ap.add_argument("-o", "--out", default="", help="输出前缀，生成 <前缀>.jsonl 和 <前缀>.txt")
    ap.add_argument("--tier", default="platform,strong",
                    help=f"保留哪些档，逗号分隔。可选 {'/'.join(TIERS)}")
    ap.add_argument("--exclude-forbidden", action="store_true",
                    help="剔掉 allow_download=false 的（抖音对 AIGC 常给 aigc_forbidden）")
    ap.add_argument("--with-jsonl", action="store_true",
                    help="连带写出 <前缀>.jsonl（体积接近源文件，默认不写）")
    ap.add_argument("--platform-needs-text", action="store_true",
                    help="platform 档要求文案/标签也有 AI 信号才算。抖音给 AI 配音、"
                         "AI 封面的普通带货视频也打 AIGC 标，露营那批 87110 条里"
                         "有 1288 条是这种，占 platform 档的一半")
    ap.add_argument("--stats", action="store_true", help="只统计不落盘")
    ap.add_argument("--samples", type=int, default=0, help="每档打印几条样例到 stderr")
    args = ap.parse_args()

    keep = {t.strip() for t in args.tier.split(",") if t.strip()}
    bad = keep - set(TIERS)
    if bad:
        ap.error(f"未知档位 {sorted(bad)}，可选 {TIERS}")

    src = Path(args.input)
    if not src.exists():
        ap.error(f"找不到 {src}")

    tiers = Counter()
    reasons = Counter()
    sigs = Counter()
    dl = Counter()
    dur_ms = 0
    kept = 0
    seen = set()
    dup = 0
    samples = {t: [] for t in TIERS}

    out_j = out_t = None
    if not args.stats:
        if not args.out:
            ap.error("要落盘就给 -o 前缀，或者加 --stats 只看统计")
        out_t = open(f"{args.out}.txt", "w", encoding="utf-8")
        out_j = open(f"{args.out}.jsonl", "w", encoding="utf-8") if args.with_jsonl else None

    total = 0
    for _, r in iter_jsonl(src):
        total += 1
        aid = str(r.get("aweme_id") or "")
        if aid in seen:
            dup += 1
            continue
        seen.add(aid)

        v = classify(r)
        t = v["tier"]
        if t == "platform" and args.platform_needs_text:
            # 平台标了 AIGC，但作品自己一个 AI 信号都没有 -> 多半是 AI 配音/封面。
            # 作者昵称带 AI（@可灵AI、@松饼大人Ai剧）也算佐证 —— 是本人在做 AI，
            # 不像 music 那条路会被 `@X创作的原声` 和 studio/Fluxxwave 污染。
            if not any(s.startswith(("desc:", "author:")) for s in v["signals"]):
                t = "none"
                reasons["platform 但无文案/作者佐证，被 --platform-needs-text 剔除"] += 1
        tiers[t] += 1
        if args.samples and len(samples[t]) < args.samples:
            samples[t].append((aid, (r.get("desc") or "").replace("\n", " ")[:90],
                               "|".join(v["signals"][:5])))
        if t == "none":
            continue
        for s in v["signals"]:
            sigs[s.split(":", 1)[0] + ":" + s.split(":")[1] if ":" in s else s] += 1

        if t not in keep:
            continue

        ai = r.get("ai") or {}
        if not ai.get("allow_download", True):
            dl[ai.get("download_fail_reason") or "unknown"] += 1
            if args.exclude_forbidden:
                reasons["被 allow_download=false 剔除"] += 1
                continue

        kept += 1
        dur_ms += int(r.get("duration_ms") or 0)
        if out_t:
            out_t.write((r.get("url") or f"https://www.douyin.com/video/{aid}") + "\n")
        if out_j:
            rec = {"aweme_id": aid, "ai_tier": t, "ai_signals": v["signals"]}
            rec.update({k: val for k, val in r.items() if k != "aweme_id"})
            out_j.write(json.dumps(rec, ensure_ascii=False, separators=(",", ":")) + "\n")

    if out_t:
        out_t.close()
    if out_j:
        out_j.close()

    w = sys.stderr
    print(f"读入 {total} 条，去重后 {len(seen)} 条（重复 {dup}）", file=w)
    print("\n档位分布:", file=w)
    for t in TIERS:
        mark = "  <- 保留" if t in keep else ""
        print(f"  {t:9s} {tiers[t]:7d}{mark}", file=w)
    ai_any = sum(tiers[t] for t in TIERS if t != "none")
    print(f"  {'AI 相关合计':9s} {ai_any:7d}  ({ai_any / max(total,1):.1%})", file=w)

    if dl:
        print("\n保留项里 allow_download=false 的原因:", file=w)
        for k, v2 in dl.most_common():
            print(f"  {k:24s} {v2:7d}", file=w)
    for k, v2 in reasons.items():
        print(f"  {k}: {v2}", file=w)

    print(f"\n最终保留 {kept} 条，总时长 {dur_ms/3600000:.2f} 小时", file=w)
    if out_t:
        print(f"写出 {args.out}.txt" + (f" 和 {args.out}.jsonl" if out_j else ""), file=w)

    if args.samples:
        for t in TIERS:
            if not samples[t]:
                continue
            print(f"\n----- {t} 样例 -----", file=w)
            for aid, d, s in samples[t]:
                print(f"  {aid}  {d}\n      信号: {s}", file=w)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())