"""元信息：标题/标签拆分，以及 records.jsonl 的追加写。

records.jsonl 是流水线唯一的落点，一条视频一行，raw 也内嵌在里面，
所以单行可能几十 KB、整个文件几百 MB。由此有两条硬性约定：

  1. 读取一律逐行流式，任何地方都不许把整个文件读进内存。
  2. `aweme_id` 必须是每行的第一个字段。续跑扫描靠 _ID_RE 从行首直接
     抠 id，不解析整行 JSON —— 400 MB 的差别是 ~1 秒 对 十几秒。
     重排字段不会报错，只会让启动悄悄变慢，所以这是契约不是排版。
"""

from __future__ import annotations

import json
import re
import threading
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

# 行首抠 id，不解析整行
_ID_RE = re.compile(r'^\{"aweme_id":\s*"(\d+)"')

# 一个标签：# 后面跟到下一个空白或下一个 # 为止。
# 抖音有 "#四川话#方言" 这种连写，所以 # 也要作为终止符。
_TAG_RE = re.compile(r"#([^\s#]+)")


# 标签两端常粘着括号、引号、逗号之类的标点，剥掉
_EDGE_PUNCT = "（）()[]【】{}「」<>《》\"'“”‘’,，.。!！?？:：;；、~～-—_/\\|"


def _clean_tag(name: str) -> str:
    """归一化一个标签候选，判定不是标签就返回空串。"""
    name = (name or "").strip().strip(_EDGE_PUNCT).strip()
    # 带 @ 的是提及，不是标签。实测 desc 里有作者手打的
    # "#(提示词出自@青木"，抖音自己也没把它收进 text_extra。
    if not name or "@" in name:
        return ""
    return name


def _dedupe(items) -> List[str]:
    """按大小写不敏感去重，保留首次出现的写法。

    抖音把 text_extra 里的标签统一转成了小写，desc 里却是作者的原始大小写。
    实测同一条视频拿到 ['一站式AI创作工作台', ..., '一站式ai创作工作台', ...]，
    同一个标签存了两份。desc 排在前面，所以保留的是原始大小写。
    """
    seen, out = set(), []
    for x in items:
        x = _clean_tag(x)
        if not x:
            continue
        k = x.casefold()
        if k not in seen:
            seen.add(k)
            out.append(x)
    return out


def split_desc(desc: Optional[str],
               text_extra: Optional[Sequence[Dict[str, Any]]] = None,
               ) -> Tuple[str, List[str]]:
    """把 desc 拆成 (纯标题, 标签列表)。

    标签以 desc 为准（正则取全集），text_extra 只用来补充。实测 12 条样本
    里有 2 条的标签在 desc 里有、text_extra 里没有，以 text_extra 为准会
    静默丢标签。

    剥标题时优先用 text_extra 给的 start/end 偏移——它能认出 "C# 教程"
    里的 # 不是标签。但偏移不一定可信（越界、错位都见过），所以每个都先
    验证 desc[start:end] 确实是 "#" + hashtag_name，验不过就退回正则。

    兜底删的是 "#" + 已知标签名 这个精确串，不是任意 #，所以 "C# 教程"
    是安全的。唯一的有损场景：标签名以 "#名字" 的形式也出现在正文里，
    那两处会一起被删掉（见 test_split_repeated_tagname_in_body_is_lossy）。
    """
    desc = (desc or "").strip()
    entries = list(text_extra or [])

    # 两件事要分开：
    #   raw_hits —— desc 里 # 后面的原文，用来剥标题（必须能字面匹配上）
    #   tags     —— 清洗+去重后的结果，是对外的字段
    # 混用会导致 "#（括号标签）" 剥不掉：清洗后的名字没有那对括号，
    # replace 匹配不上，标签就留在标题里了。
    raw_hits = _TAG_RE.findall(desc)
    from_te = [e.get("hashtag_name") for e in entries if e.get("hashtag_name")]
    tags = _dedupe(raw_hits + from_te)

    if not desc:
        return "", tags

    # 能用偏移剥掉的标签，逐个记下要挖掉的区间
    spans: List[Tuple[int, int]] = []
    covered = set()
    for e in entries:
        name = e.get("hashtag_name")
        start, end = e.get("start"), e.get("end")
        if not name or not isinstance(start, int) or not isinstance(end, int):
            continue
        if not (0 <= start < end <= len(desc)):
            continue
        if desc[start:end] != "#" + name:
            continue
        spans.append((start, end))
        covered.add(name)

    title = desc
    for start, end in sorted(spans, reverse=True):   # 从后往前挖，偏移不会串位
        title = title[:start] + title[end:]

    # 偏移没盖到的，用 desc 里的原文删（不是清洗过的 tags）
    for hit in raw_hits:
        if hit not in covered:
            title = title.replace("#" + hit, "")

    return re.sub(r"\s+", " ", title).strip(), tags


def build_record(aweme, *, duration_sec: float, cos_mp4: str, cos_mp3: str,
                 size_mp4: int, size_mp3: int,
                 source_url: str = "") -> Dict[str, Any]:
    """组装一条 record。字段顺序是有意的，aweme_id 必须在最前。"""
    raw = aweme.raw or {}
    title, tags = split_desc(aweme.desc, raw.get("text_extra"))
    author = raw.get("author") or {}
    stat = raw.get("statistics") or {}
    music = raw.get("music") or {}
    video = raw.get("video") or {}
    ct = int(aweme.create_time or 0)

    return {
        "aweme_id": str(aweme.aweme_id),
        "title": title,
        "desc": aweme.desc,
        "tags": tags,
        "author": {
            "nickname": aweme.author,
            "uid": str(author.get("uid") or ""),
            "sec_uid": str(author.get("sec_uid") or ""),
        },
        "create_time": ct,
        "create_time_str": (time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(ct))
                            if ct else ""),
        "duration_sec": duration_sec,
        # 服务端声明值，与 ffprobe 解出来的不等（实测 41750 vs 41703 ms）。
        # 两个都留着，切分要用 duration_sec。
        "duration_ms_api": video.get("duration") or raw.get("duration") or 0,
        "stats": {
            "digg": stat.get("digg_count", 0),
            "comment": stat.get("comment_count", 0),
            "collect": stat.get("collect_count", 0),
            "share": stat.get("share_count", 0),
            "play": stat.get("play_count", 0),
        },
        "music": {
            "title": (music.get("title") or "").strip(),
            "author": (music.get("author") or "").strip(),
            "url": aweme.music_url or "",
        },
        "cover_url": aweme.cover_urls[0] if aweme.cover_urls else "",
        "video": {
            "uri": aweme.uri,
            "width": video.get("width") or 0,
            "height": video.get("height") or 0,
        },
        "cos_mp4": cos_mp4,
        "cos_mp3": cos_mp3,
        "size_mp4": size_mp4,
        "size_mp3": size_mp3,
        "source_url": source_url or f"https://www.douyin.com/video/{aweme.aweme_id}",
        "fetched_at": int(time.time()),
        "raw": raw,
    }


class RecordLog:
    """records.jsonl 的追加写 + 已完成 id 集合。多线程安全。"""

    def __init__(self, path: Path) -> None:
        self.path = Path(path)
        self._lock = threading.Lock()
        self._ids = self._scan()

    def _scan(self) -> set:
        """流式扫已有 id。只看行首，不解析整行。"""
        ids = set()
        if not self.path.exists():
            return ids
        with open(self.path, "r", encoding="utf-8", errors="replace") as f:
            for line in f:
                m = _ID_RE.match(line)
                if m:
                    ids.add(m.group(1))
        return ids

    def __contains__(self, aweme_id: str) -> bool:
        return str(aweme_id) in self._ids

    def __len__(self) -> int:
        return len(self._ids)

    @property
    def ids(self) -> set:
        return set(self._ids)

    def _needs_newline(self) -> bool:
        """上一次可能被 Ctrl-C 打断，留下没换行的半行。"""
        try:
            if self.path.stat().st_size == 0:
                return False
        except OSError:
            return False
        with open(self.path, "rb") as f:
            f.seek(-1, 2)
            return f.read(1) != b"\n"

    def write(self, record: Dict[str, Any]) -> bool:
        """追加一条。已存在同 id 则跳过，返回 False。"""
        aweme_id = str(record.get("aweme_id") or "")
        if not aweme_id:
            raise ValueError("record 缺 aweme_id")

        with self._lock:
            if aweme_id in self._ids:
                return False
            self.path.parent.mkdir(parents=True, exist_ok=True)
            prefix = "\n" if self._needs_newline() else ""
            line = json.dumps(record, ensure_ascii=False, separators=(",", ":"))
            if not line.startswith('{"aweme_id":'):
                raise ValueError("aweme_id 必须是第一个字段，见本模块开头的约定")
            with open(self.path, "a", encoding="utf-8") as f:
                f.write(prefix + line + "\n")
            self._ids.add(aweme_id)
            return True


def iter_records(path: Path):
    """逐行读 records.jsonl，跳过坏行（可能有被打断留下的半行）。"""
    path = Path(path)
    if not path.exists():
        return
    with open(path, "r", encoding="utf-8", errors="replace") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                yield json.loads(line)
            except ValueError:
                continue