#!/usr/bin/env python
"""把 records.jsonl 导成交付用的样本 xlsx。

列是甲方给的字段表（序号跟着他们的编号走，中间 4/5/9 他们没要）：

    1 样本 ID          必填   aweme_id
    2 视频标题          必填   title，为空时退回 desc 首行
    3 平台             必填   固定「抖音」
    6 生成方式                AI 生成 / AI 辅助 / 疑似 AI / 未标注
    7 生成工具 / 模型          从文案标签里认出来的工具名，逗号分隔
    8 题材 / 内容类别          按标签归类
   10 分辨率 / 清晰度    必填   1024x576（540p）
   11 发布时间          必填   2026-08-06 14:12:47
   12 是否有 AI 标识/水印 必填   抖音 risk_infos 打的标，原文照抄

用法:
    python export_sample_xlsx.py                      # 全量
    python export_sample_xlsx.py --ai-only            # 只要判定为 AI 的
    python export_sample_xlsx.py --limit 500 -o 样本.xlsx
"""

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 re
import sys
from pathlib import Path

from douyin.aiclass import _fields, _scan, classify

# risk_infos.type 里属于 AIGC 的那几个码，别的都是跟 AI 无关的免责声明
# （45=虚构演绎、336/351=个人观点、54=来源于网络…）。见 README §4.5。
_AIGC_TYPES = {44, 47, 76, 328}

# 生成方式：抖音标签的措辞本身就分了轻重，直接用它，别自己发明
_HOW_BY_TYPE = {
    44: "AI 生成（作者声明）",
    76: "AI 生成（平台判定）",
    328: "AI 生成（平台判定）",
    47: "疑似 AI 生成（平台检测）",
}

# 题材归类。**顺序即优先级**，第一个命中的赢 —— 一条片子既是「漫剧」又带
# 「剧情」很常见，先匹配到的那个更具体。
#
# 词要挑「只在这个题材里出现」的。踩过的坑：「古风/国风/画质」放进
# 写真/摄影，结果把一整个仙侠短片系列（169/300 条）全吸过去了 ——
# 那些是叙事片，只是标签带古风。宽词一律不要，宁可落进「其他」。
_GENRES = [
    ("漫剧 / 动画", "漫剧|动漫|动画|二次元|国漫|番剧|漫画"),
    ("短剧 / 剧情", "短剧|剧情|微短剧|电影|影视|短片|反转|结局|"
                    "第[一二三四五六七八九十\\d]+集|合集"),
    ("教程 / 科普", "教程|教学|技巧|入门|保姆级|科普|干货|工作流|提示词"),
    ("资讯 / 新闻", "新闻|资讯|报道|快讯|媒体精选|发布会|财经|股市"),
    ("带货 / 电商", "好物|同款|下单|优惠|源头工厂|旗舰店|带货|种草|清仓"),
    ("户外 / 露营", "露营|帐篷|户外|野营|钓鱼|自驾|徒步|登山"),
    ("音乐 / 舞蹈", "音乐|歌曲|翻唱|舞蹈|跳舞"),
    ("写真 / 摄影", "写真|摄影|人像|壁纸|头像"),
    ("游戏", "手游|端游|原神|王者荣耀|电竞|无畏契约|崩坏"),
    ("生活 / 日常", "日常|vlog|美食|旅行|萌宠|搞笑"),
]
_GENRE_RE = [(name, re.compile(pat, re.I)) for name, pat in _GENRES]


def _risk(rec: dict) -> tuple:
    """(type, content)。抖音把 AI 标识放在 raw.risk_infos 里。"""
    ri = (rec.get("raw") or {}).get("risk_infos") or {}
    if not isinstance(ri, dict):
        return None, ""
    return ri.get("type"), (ri.get("content") or "").strip()


def how_generated(rec: dict, verdict: dict) -> str:
    t, _ = _risk(rec)
    if t in _HOW_BY_TYPE:
        return _HOW_BY_TYPE[t]
    if verdict["tier"] == "strong":
        return "AI 生成（作品自述）"
    if verdict["tier"] in ("topic", "author"):
        return "疑似 AI 相关"
    return "未标注"


def tools_used(verdict: dict) -> str:
    """从判定信号里把工具名摘出来。信号形如 `desc:工具:即梦`。"""
    names = []
    for s in verdict["signals"]:
        parts = s.split(":")
        if len(parts) >= 3 and parts[1] in ("工具", "漫剧"):
            n = parts[2]
            if n not in names:
                names.append(n)
    return "、".join(names)


def genre(rec: dict) -> str:
    text = " ".join([rec.get("title") or "", rec.get("desc") or ""]
                    + list(rec.get("tags") or []))
    for name, rx in _GENRE_RE:
        if rx.search(text):
            return name
    return "其他"


def resolution(rec: dict) -> str:
    v = rec.get("video") or {}
    w, h = v.get("width"), v.get("height")
    ratio = ((rec.get("raw") or {}).get("video") or {}).get("ratio") or ""
    if w and h:
        return f"{w}x{h}" + (f"（{ratio}）" if ratio else "")
    return ratio or ""


def ai_mark(rec: dict) -> str:
    """是否有 AI 标识/水印 —— 抖音的原文照抄，没有就写「无」。
    只认 _AIGC_TYPES；45/336/351 那些是虚构演绎、个人观点，跟 AI 无关。"""
    t, content = _risk(rec)
    if t in _AIGC_TYPES and content:
        return f"有：{content}"
    return "无"


def title_of(rec: dict) -> str:
    """标题 + 标签拼一格。

    records.jsonl 里的 title 是 split_desc 把标签摘掉之后的正文，标签单独放
    在 tags 里。甲方只要一列，所以拼回去 —— 顺序按原文案的习惯：正文在前，
    `#标签` 跟在后面。
    """
    t = (rec.get("title") or "").strip()
    if not t:
        d = (rec.get("desc") or "").strip().splitlines()
        t = d[0][:120] if d else ""
    tags = " ".join(f"#{x}" for x in (rec.get("tags") or []) if x)
    return (t + " " + tags).strip() if tags else t


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description="records.jsonl -> 交付样本 xlsx")
    ap.add_argument("-i", "--input", default="records/records.jsonl")
    ap.add_argument("-o", "--out", default="exports/AI视频样本.xlsx")
    ap.add_argument("--ai-only", action="store_true",
                    help="只导出判定为 AI 的（platform/strong 档）")
    ap.add_argument("--limit", type=int, default=0)
    args = ap.parse_args(argv)

    try:
        from openpyxl import Workbook
    except ImportError:
        print("需要 openpyxl：pip install openpyxl", file=sys.stderr)
        return 2

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

    # write_only：11 万行常规模式会把内存吃光
    wb = Workbook(write_only=True)
    ws = wb.create_sheet("样本")
    # 前 9 列是甲方字段表要的，后面几列是核对/取数用的：题材判得对不对，
    # 光看标题看不出来，得有标签和链接才能点开验；COS 路径是拿原始文件的地址。
    # 作者插在「平台」后面 —— 甲方编号 1/2/3 的三列保持原位原序，
    # 作者正好落进他们跳过的 4/5 那两个空号里。
    ws.append(["样本 ID", "视频标题", "平台", "作者",
               "生成方式", "生成工具 / 模型",
               "题材 / 内容类别", "分辨率 / 清晰度", "发布时间",
               "是否有 AI 标识 / 水印",
               "时长(秒)", "视频链接",
               "COS 视频路径", "COS 音频路径"])

    n = written = bad = 0
    with open(src, "r", encoding="utf-8", errors="replace") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            n += 1
            try:
                rec = json.loads(line)
            except ValueError:
                bad += 1
                continue
            v = classify(rec)
            if args.ai_only and v["tier"] not in ("platform", "strong"):
                continue
            # 每格都过一遍 str()：openpyxl 会把 None 原样写成 "None"
            ws.append([
                str(rec.get("aweme_id") or ""),
                title_of(rec),
                "抖音",
                (rec.get("author") or {}).get("nickname") or "",
                # 生成方式 / 生成工具：留空，人工填。
                # 自动推的那版（how_generated / tools_used）只能看抖音的标和
                # 文案里写没写工具名，作者不写就抓不到 —— 实测 44.8% 的工具
                # 列是空的，而"生成方式"里 81.8% 全都压在"作者声明"一个值上，
                # 区分度太低，不如空着让人填。两个函数留着，要用加回来即可。
                "",
                "",
                genre(rec),
                resolution(rec),
                rec.get("create_time_str") or "",
                ai_mark(rec),
                round(float(rec.get("duration_sec") or 0), 1),
                rec.get("source_url")
                or f"https://www.douyin.com/video/{rec.get('aweme_id')}",
                rec.get("cos_mp4") or "",
                rec.get("cos_mp3") or "",
            ])
            written += 1
            if args.limit and written >= args.limit:
                break
            if written % 20000 == 0:
                print(f"  已写 {written} 行…", flush=True)

    wb.save(args.out)
    print(f"读入 {n} 行（坏行 {bad}），写出 {written} 行 -> {args.out}")
    return 0


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