"""LLM 字幕过滤节点：对 OCR 识别出的 SRT 做二次过滤。

两级判断，**默认只跑规则层**（use_llm=0；LLM 层误删真实对话的代价过高，
依据见 docs/decisions.md）：

1. **确定性规则层**（不调 LLM）：横线装饰、HTML/水印 token、URL/邮箱、
   单双 ASCII 字符等 OCR 噪声直接删除。模式稳定可判，省 token 且结果确定。

2. **LLM 五类分类层**（use_llm=1 时启用）：把目标字幕连同前后各
   context_size 条纯文本交给 LLM，输出五个类别之一：
   - garbage（乱码/残缺/装饰）、overlay（水印/播放器覆盖层）、
     noise（无实义杂项）→ 删除；
   - repeat（内容性重复，如语气词/呻吟）、dialogue（正常对话）→ 保留。
   未识别输出一律回退 dialogue（宁滥勿缺）。

   该层的两个配套机制：
   - **文本去重**：相同文本（忽略空白与大小写）只调一次 LLM，结果复用，
     保证同一句话判定一致并减少调用量；
   - **上下文净化**：喂给 LLM 的上下文先剔除规则层已判定的垃圾，避免覆盖层
     噪声污染场景判断而误删真实对话。
"""

from __future__ import annotations

import json
import os
import re
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path

from nodes.adaptive_pool import AdaptiveThreadPool
from wov_app.logging import get_logger
from wov_sdk.models import InvokeRequest, InvokeResponse

logger = get_logger("llm-filter")

# 匹配 SRT 条目：时间轴行 + 文本（文本可多行），到下一个序号行或文末结束。
_SRT_BLOCK_RE = re.compile(
    r"(\d{2}:\d{2}:\d{2},\d{3})\s*-->\s*(\d{2}:\d{2}:\d{2},\d{3})\s*\n(.*?)(?=\n\s*\d+\s*\n|\Z)",
    re.DOTALL,
)

# 目标字幕标记：提示词用该标记指明需要判断的那一条字幕。
TARGET_MARK = "【目标】"

# 默认上下文窗口：目标字幕前后各取 10 条。
DEFAULT_CONTEXT_SIZE = 10

# LLM 输出类别：garbage/overlay/noise 删除；repeat/dialogue 保留。
CATEGORY_GARBAGE = "garbage"
CATEGORY_OVERLAY = "overlay"
CATEGORY_NOISE = "noise"
CATEGORY_REPEAT = "repeat"
CATEGORY_DIALOGUE = "dialogue"
_ALL_CATEGORIES = (
    CATEGORY_GARBAGE,
    CATEGORY_OVERLAY,
    CATEGORY_NOISE,
    CATEGORY_REPEAT,
    CATEGORY_DIALOGUE,
)
DELETE_CATEGORIES = {CATEGORY_GARBAGE, CATEGORY_OVERLAY, CATEGORY_NOISE}

# 判定存档文件名：位于节点 output_dir，每行 {"index": 条目标引, "category": 类别}。
# 每条 LLM 判定成功即追加一行；进程被杀/节点失败（如 429 限流）后重跑时，
# 只对未判定的条目重新调用 LLM，已判定结果直接复用（类似 OCR 的断点存档）。
_PARTIAL_NAME = "filter_partial.jsonl"

# 判定存档追加写锁：多线程判定并发完成时串行化追加，避免行交错。
_partial_lock = threading.Lock()


def _load_partial(output_dir: Path) -> dict[int, str]:
    """读取判定存档，返回 {条目标引: 类别}；无存档/损坏行跳过。"""
    path = output_dir / _PARTIAL_NAME
    if not path.is_file():
        return {}
    result: dict[int, str] = {}
    for line in path.read_text(encoding="utf-8").splitlines():
        if not line.strip():
            continue
        try:
            item = json.loads(line)
        except json.JSONDecodeError:
            # 进程被杀时可能残留半行写入：跳过，对应条目视为未判定。
            continue
        result[int(item["index"])] = str(item["category"])
    return result


def _append_partial(output_dir: Path, index: int, category: str) -> None:
    """线程安全地把一条判定结果追加到存档（成功判定后立即落盘）。"""
    with _partial_lock:
        with (output_dir / _PARTIAL_NAME).open("a", encoding="utf-8") as fh:
            fh.write(json.dumps({"index": index, "category": category}, ensure_ascii=False) + "\n")

# 规则层正则：横线装饰（含全角/半角横线、下划线、中点、句点等符号组合）。
_DASH_RE = re.compile(r"^[\s\-—_~=•・。．.、]+$")
# 规则层正则：URL / 邮箱。
_URL_OR_MAIL_RE = re.compile(r"^(https?://|www\.)\S+$|^[\w.+-]+@[\w.-]+\.\w+$")
# 规则层正则：裸网址/域名（含中文夹杂的注册地址，如 "水火地址 489155.com"）。
_DOMAIN_RE = re.compile(
    r"[\w-]+\.(?:com|net|org|cn|tv|me|io|xyz|cc|top|info|biz)(?:[/\s.,;:：!?）)]|$)",
    re.IGNORECASE,
)
# 规则层正则：HTML/脚本/播放器水印模式（OCR 常把网页界面识别成这类文本）。
_HTML_MARK_RE = re.compile(
    r"html\s*code|<\s*[a-z][^>]*>|javascript|web\s*address|watermark|sign\s*in",
    re.IGNORECASE,
)
# 规则层正则：播放器/作品编号水印（VLM 常把画面角落的编号识别成短串）。
# 覆盖水印编号形态，如 SPHO-1 / PHO一号馆 / NO.1专用 / PHD-手術 / SP10-1型。
# 限定为**不含汉字的编号形态**（或纯形态串），避免误伤正常英文对白。
_SERIAL_MARK_RE = re.compile(
    r"^(?:SPH|SPHO|SPIO|SPNO|SP10|PHO|PH0|PHD|P10|NO\.|SP\s*\d)"
    r"[\w\s.+#\-一-鿿]{0,12}$",
    re.IGNORECASE,
)
# 规则层正则：日期/时间戳（OCR 把画面日期当成字幕）。
_DATE_ONLY_RE = re.compile(r"^\d{4}\s*[-/年]\s*\d{1,2}\s*[-/月]\s*\d{1,2}\s*日?$")
# 规则层正则：VLM 提示词回显（glm-ocr 偶尔把系统提示当作识别结果输出）。
_VLM_ECHO_RE = re.compile(
    r"no text is visible|image is blurry|does not contain (?:any )?text|no text visible",
    re.IGNORECASE,
)
# 规则层正则：演员/出演标注（片头片尾覆盖层）。**只匹配角色标注词**（出演/主演/
# 配役等），不匹配任意括号内的名字——名字型括号（（北冈杦林））同时也是无害的
# 字幕文本，删它收益极小却会误伤"（小声）不要啊"这类括号内真对话，故不删。
_CAST_MARK_RE = re.compile(r"^[（(]\s*(?:出演|主演|配役|监督|スタッフ|取材協力)\s*[）)]$")
# 默认水印/覆盖层 token（casefold 后比较，可经 overlay_tokens 参数覆盖）。
DEFAULT_OVERLAY_TOKENS = frozenset(
    {"html", "background", "___", "cleaning", "buffering", "loading", "marketing"}
)
# 长文本保护阈值：≥ 该长度的文本，noise 类别不构成删除依据（LLM 会把完整
# 对话句误判为 noise，长度是兜底手段，只用于保护不用于删除）。
DEFAULT_MIN_KEEP_LEN = 12

# 纯数字/小数点组合（"4.0"/"2.0"/"10-1期B" 中的纯数值形态）。
_NUMBER_ONLY_RE = re.compile(r"^\d{1,3}[.,]\d{1,2}$")

# LLM 分类层开关：默认关闭，只跑确定性规则层（原因见模块 docstring）。
DEFAULT_USE_LLM = False


def _has_cjk(text: str) -> bool:
    """是否含 CJK 汉字：单字"嗯/好"等可能是内容，规则层不直接删。"""
    return any("\u4e00" <= ch <= "\u9fff" for ch in text)


def parse_srt(text: str) -> list[dict]:
    """解析 SRT 文本为条目列表：[{"start", "end", "text"}]。"""
    entries: list[dict] = []
    for match in _SRT_BLOCK_RE.finditer(text):
        entries.append(
            {
                "start": match.group(1),
                "end": match.group(2),
                "text": match.group(3).strip(),
            }
        )
    return entries


def serialize_srt(entries: list[dict]) -> str:
    """把条目列表序列化为标准 SRT 文本（序号重新从 1 编号）。"""
    blocks = [
        f"{index}\n{entry['start']} --> {entry['end']}\n{entry['text']}"
        for index, entry in enumerate(entries, start=1)
    ]
    return "\n\n".join(blocks) + "\n"


def _rule_verdict(text: str, overlay_tokens: set[str]) -> bool | None:
    """确定性规则层：返回 True（删除）/ None（交给 LLM 多维判断）。

    规则覆盖 OCR 噪声的稳定模式：空文本、纯横线装饰、URL/邮箱、裸网址
    域名、HTML/水印模式（html code/标签/javascript 等）、水印 token、单双
    ASCII 字符。含 CJK 的短文本不算垃圾，因为"嗯/好"等可能是内容；
    其余情况返回 None 交由 LLM 分类。
    """
    t = text.strip()
    if not t:
        return True
    if _DASH_RE.match(t):
        return True
    # 日期/编号型：先于域名判断（避免 "2011-11-27" 被当作普通文本）。
    if _DATE_ONLY_RE.match(t):
        return True
    if _SERIAL_MARK_RE.match(t):
        return True
    if _VLM_ECHO_RE.search(t):
        return True
    if _URL_OR_MAIL_RE.match(t) or _DOMAIN_RE.search(t):
        return True
    if _HTML_MARK_RE.search(t):
        return True
    if t.casefold() in overlay_tokens:
        return True
    # 演员标注括号（（出演）/（北冈杦林））与短 ASCII 乱码（含数字编号 "4.0"）。
    # 含汉字的括号内容（如"（小声）不要啊"）天然不命中 _CAST_MARK_RE 的字符集。
    if _CAST_MARK_RE.match(t):
        return True
    if len(t) <= 2 and not _has_cjk(t):
        return True
    # 纯数字/小数点组合（"4.0"、"2.0"）：OCR 把画面数值识别成条目。
    if _NUMBER_ONLY_RE.match(t):
        return True
    return None


def _dedup_key(text: str) -> str:
    """文本去重键：去掉全部空白并统一大小写。

    OCR 同一句字幕常带/不带空格（"可没法胜任" vs "可 没法胜任"），
    视为同一文本以保证判定一致；也用于把 LLM 调用量降到唯一文本数。
    """
    return "".join(text.split()).casefold()


def _should_delete(category: str, text: str, min_keep_len: int) -> bool:
    """按 LLM 类别与长文本保护决定是否删除。

    repeat/dialogue 一律保留；garbage/overlay 一律删除（明确的垃圾信号）；
    noise 仅对短文本删除：LLM 会把完整对话句误判为 noise，≥min_keep_len 的
    长文本一律保留。长度只用于保护，不用于删除。
    """
    if category not in DELETE_CATEGORIES:
        return False
    if category == CATEGORY_NOISE and len(text) >= min_keep_len:
        return False
    return True

def _judge_category(
    entries: list[dict], index: int, context_size: int, params: dict,
    overlay_tokens: set[str] | None = None,
) -> str:
    """调用 LLM 把目标字幕归入五类之一，返回类别词（未识别回退 dialogue）。

    请求体只含目标字幕及其前后各 context_size 条字幕的纯文本（无时间戳），
    目标字幕用 TARGET_MARK 标记；模型只输出一个类别英文单词。
    输出无法识别（空/乱码/旧式"保留"）时回退 dialogue，宁滥勿缺。
    """
    start = max(0, index - context_size)
    end = min(len(entries), index + context_size + 1)
    target_pos = index - start
    # 上下文净化：喂给 LLM 的是**过滤后的字幕**——相邻条目若被确定性规则层
    # 识别为垃圾（横线/HTML/水印 token/URL/裸域名等）直接从上下文中剔除，
    # 避免覆盖层垃圾污染 LLM 对整段场景的判断（误删相邻的真实对话）。
    tokens = overlay_tokens if overlay_tokens is not None else DEFAULT_OVERLAY_TOKENS
    lines = []
    for pos, entry in enumerate(entries[start:end]):
        if pos != target_pos and _rule_verdict(entry["text"], tokens) is True:
            continue
        text = entry["text"]
        lines.append(f"{TARGET_MARK}{text}" if pos == target_pos else text)

    # LLM 兼容接口配置：地址/Key/模型/超时均可通过环境变量覆盖（默认 SiliconFlow）。
    api_base = os.getenv(
        "LLM_API_BASE",
        "https://api.siliconflow.cn/v1/chat/completions",
    )
    api_key = os.getenv("LLM_API_KEY", "")
    request_timeout = float(os.getenv("LLM_TIMEOUT_SECONDS", "60"))
    model = str(params.get("model") or os.getenv("LLM_MODEL", "Qwen/Qwen3.5-35B-A3B"))
    system_prompt = (
        "你是字幕质量过滤器。用户会提供一段字幕序列（纯文本，不含时间戳），"
        f"其中用{TARGET_MARK}标记的字幕是需要判断的目标。"
        "请把目标字幕归入以下五个类别之一：\n"
        "garbage：垃圾字符（乱码、残缺、装饰性符号、横线）\n"
        "overlay：水印、网页/播放器/字幕组等覆盖层文本，不是视频对白\n"
        "noise：与上下文无关、无实际语义的杂项\n"
        "repeat：内容性重复（如语气词、呻吟、重复的感叹或对话），属于内容本身\n"
        "dialogue：正常对话\n"
        "字幕序列已经过确定性规则过滤（装饰性横线、HTML、网址/水印等已被剔除），"
        "请只依据剩下的对话内容判断目标字幕，不要臆测被过滤掉的部分。\n"
        "只输出一个类别英文单词，不要输出其他内容。"
    )
    body = {
        "model": model,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": "\n".join(lines)},
        ],
        # 关闭推理模式：Qwen3 等模型默认会把思考过程写入 reasoning_content，
        # 导致 content 为空或包含多余内容。
        "enable_thinking": False,
        # 类别词很短，输出上限给得很小即可。
        "max_tokens": 16,
    }
    headers = {"Content-Type": "application/json"}
    # 配置了 Key 时附带 Bearer 鉴权头。
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    request = urllib.request.Request(
        api_base,
        data=json.dumps(body).encode("utf-8"),
        headers=headers,
        method="POST",
    )
    # 429（限流）与 5xx（服务端错误）时指数退避重试：请求被限流时让出时间，
    # 使窗口平均响应变慢，触发自适应线程池"慢响应减线程"，并发自动回落到
    # 限流配额内；最多尝试 3 次，耗尽仍失败则抛出，由调用方（任务）重新处理。
    max_attempts = 3
    retry_delay = 1.0
    for attempt in range(max_attempts):
        try:
            with urllib.request.urlopen(request, timeout=request_timeout) as response:
                payload = json.loads(response.read().decode("utf-8"))
            break
        except urllib.error.HTTPError as exc:
            if exc.code != 429 and not (500 <= exc.code < 600):
                raise
            if attempt >= max_attempts - 1:
                raise
            time.sleep(retry_delay)
            retry_delay *= 2
    content = str(payload["choices"][0]["message"]["content"]).strip().lower()
    # 精确匹配五个类别词（兼容"garbagexxx"这类多余输出）。
    for category in _ALL_CATEGORIES:
        if content == category or content.startswith(category):
            return category
    # 旧式"删除/保留"回答兼容：含"删除"视为垃圾，其余一律保留。
    if "删除" in content:
        return CATEGORY_GARBAGE
    return CATEGORY_DIALOGUE


def invoke(request: InvokeRequest) -> InvokeResponse:
    """过滤 SRT 中多余/无意义的字幕，产物为 filtered.srt。

    流程：规则层（确定性删除）→ LLM 层（去重后按唯一文本多维分类，
    上下文为过滤后的字幕）→ 保留条重新编号输出。
    """
    srt_uri = request.inputs.get("srt_uri")
    if not srt_uri:
        return InvokeResponse(status="failed", error="srt_uri is required")
    srt_path = Path(srt_uri)
    if not srt_path.is_file():
        return InvokeResponse(status="failed", error="srt file not found")

    entries = parse_srt(srt_path.read_text(encoding="utf-8"))
    context_size = int(request.params.get("context_size", DEFAULT_CONTEXT_SIZE))
    min_keep_len = int(request.params.get("min_keep_len", DEFAULT_MIN_KEEP_LEN))
    # 水印 token 可经参数覆盖（JSON 数组字符串或列表），默认内置常见覆盖层词。
    raw_tokens = request.params.get("overlay_tokens")
    if isinstance(raw_tokens, str) and raw_tokens.strip():
        raw_tokens = json.loads(raw_tokens)
    overlay_tokens = {
        str(t).casefold() for t in (raw_tokens or DEFAULT_OVERLAY_TOKENS)
    }
    # 去重开关：默认开；关掉时每个条目独立调用 LLM（不省调用，判定各自独立）。
    dedupe = str(request.params.get("dedupe", "1")) not in ("0", "false", "False")
    # LLM 分类层开关：默认关闭，只跑确定性规则层（宁多留不漏删）；关闭原因
    # 见模块 docstring 与 docs/decisions.md。需要该层时用 params.use_llm=1。
    use_llm = str(request.params.get("use_llm", "1" if DEFAULT_USE_LLM else "0")) not in (
        "0", "false", "False",
    )

    # 阶段 1：确定性规则层（不调 LLM）。
    rule_verdicts = [_rule_verdict(entry["text"], overlay_tokens) for entry in entries]
    llm_needed = [i for i, verdict in enumerate(rule_verdicts) if verdict is None] if use_llm else []

    # 阶段 2：LLM 分类层（去重：相同文本只判一次，上下文取首次出现）。
    # 节点级断点存档：output_dir 提前建好，重跑时只重判未判定条目。
    output_dir = Path(request.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    partial = _load_partial(output_dir)
    cat_by_index: dict[int, str] = dict(partial)
    if llm_needed:
        if dedupe:
            first_of_key: dict[str, int] = {}
            pool_indices: list[int] = []
            for i in llm_needed:
                key = _dedup_key(entries[i]["text"])
                if key not in first_of_key:
                    first_of_key[key] = i
                    # 断点续跑：该键首次出现已在存档判定过则跳过（结果复用）。
                    if i not in partial:
                        pool_indices.append(i)
        else:
            # 断点续跑：只处理未判定的条目。
            pool_indices = [i for i in llm_needed if i not in partial]
        # 单条判断的工作函数：返回类别词；overlay_tokens 用于上下文净化。
        def judge_one(index: int) -> str:
            try:
                category = _judge_category(entries, index, context_size, request.params, overlay_tokens)
            except urllib.error.HTTPError as exc:
                # 限流/服务端错误：通知线程池临时降低最大并发，避免持续超配额。
                if exc.code == 429 or 500 <= exc.code < 600:
                    pool.report_failure()
                raise
            # 判定成功立即落盘（断点存档）：失败/中断后重跑不重复调用已判定条目。
            _append_partial(output_dir, index, category)
            return category
        # 进度日志：打印已判定条数、总数、平均处理速度（条/s）、最近窗口
        # 平均单条耗时与当前线程数（与 OCR 节点同一回调协议）。
        def log_progress(done: int, total: int, rate: float, avg_time: float, workers: int) -> None:
            logger.info(
                "字幕判定进度 %d/%d 条 (%.1f 条/s, 平均 %.2fs/条, 线程 %d/%d)",
                done, total, rate, avg_time, workers, pool.max_workers,
            )

        # 自适应并发调用 LLM：按响应快慢弹性伸缩，避免压垮接口。
        pool = AdaptiveThreadPool(
            worker=judge_one,
            on_progress=log_progress,
            min_workers=int(request.params.get("pool_min_workers", 1)),
            max_workers=int(request.params.get("pool_max_workers", 16)),
            window_seconds=float(request.params.get("pool_window_seconds", 10.0)),
            fast_threshold=float(request.params.get("pool_fast_threshold", 0.3)),
            slow_threshold=float(request.params.get("pool_slow_threshold", 1.0)),
        )
        categories = pool.map(pool_indices)

        # 限流/服务端错误已在 worker 内 report_failure 临时降并发：把失败的
        # 条目在收紧后的并发下重试一轮（_judge_category 内还有 429/5xx 退避），
        # 二次仍失败才整体失败——避免 20 并发一拥而上被限流打挂整个任务。
        failed = [
            index for index, category in zip(pool_indices, categories)
            if isinstance(category, Exception)
        ]
        if failed:
            logger.info("判定失败 %d 条，降并发后重试", len(failed))
            retried = pool.map(failed)
            for index, category in zip(failed, retried):
                if isinstance(category, Exception):
                    # 二次仍失败（限流持续/非限流错误）：整体失败，由任务重试。
                    return InvokeResponse(status="failed", error=str(category))
                cat_by_index[index] = category
        for index, category in zip(pool_indices, categories):
            # 首次失败的条目已在重试分支处理，这里只登记成功结果。
            if isinstance(category, Exception):
                continue
            cat_by_index[index] = category
        if dedupe:
            # 去重填充：与首次出现同键的条目复用同一类别，保证判定一致。
            for i in llm_needed:
                if i not in cat_by_index:
                    cat_by_index[i] = cat_by_index[first_of_key[_dedup_key(entries[i]["text"])]]

    # 阶段 3：合并规则与 LLM 判定，输出保留条。
    kept: list[dict] = []
    removed = 0
    for i, (entry, rule_verdict) in enumerate(zip(entries, rule_verdicts)):
        # 规则层命中即删（True）；未命中则按 LLM 类别判定。
        if rule_verdict is True:
            removed += 1
            continue
        # LLM 层关闭时，规则层未命中的条目一律保留（不做分类判定）。
        if not use_llm:
            kept.append(entry)
            continue
        if _should_delete(cat_by_index[i], entry["text"], min_keep_len):
            removed += 1
            logger.info("删除无意义字幕 %d: %r", i + 1, entry["text"][:40])
        else:
            kept.append(entry)

    output_dir = Path(request.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / "filtered.srt"
    output_path.write_text(serialize_srt(kept), encoding="utf-8")
    logger.info("字幕过滤完成: 保留 %d 条, 删除 %d 条", len(kept), removed)
    return InvokeResponse(
        status="completed",
        outputs={"srt_uri": str(output_path), "kept": len(kept), "removed": removed},
    )