"""LLM 翻译节点。

单体版中作为进程内节点模块，由调度器直接调用。接收 SRT，提取纯文本行
分批调用 LLM，再把译文回填到原 SRT 结构并输出 cn.srt。

关键修复（见 tests/test_translation_line_alignment.py）：

1. **提示词强化**：要求"逐行独立翻译 + 碎片句按语境独立成行 + 禁止合并/拆分"，
   从源头减少 LLM 因语义碎片而重排断句、导致行数不一致。

2. **行数对齐（_repair_batch）**：LLM 偶发多拆/少拆一行会让后续所有字幕文本
   相对时间戳整体错位（时间戳从原文复制、文本却错贴到其他时间——程序按时戳
   看不出问题，实测 run_51242078d76e 大量批次出现 21/19 行 vs 输入 20 行）。
   处理：多行 -> 末尾多余行合并到前一行；少行 -> 重试该批（内容缺失无法靠
   占位恢复），仍不足则补空串占位（宁缺勿错位）。

3. **system_prompt 拼接 bug**：圆括号内一旦出现 f-string 赋值（表达式），
   隐式字符串拼接失效，整体变成 tuple；json 序列化后发出去的 content 是数组，
   API 返回 400 invalid parameter。必须用 + 显式拼接为单个字符串。
"""

from __future__ import annotations

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

from wov_app.logging import get_logger

from wov_sdk.models import InvokeRequest, InvokeResponse
from nodes.subtitle_cleanup import clean_srt_text
from nodes.proper_nouns import build_proper_noun_rule
# 单次 LLM 请求携带的字幕行数；过大会超出模型上下文，过小则请求次数过多。
CHUNK_SIZE = 20

# 批次翻译重试次数（LLM 偶发少行时重发本批，内容缺失无法靠占位恢复）。
MAX_BATCH_RETRIES = 3

# 节点运行日志：翻译分批进度与处理速度输出到主进程控制台。
logger = get_logger("llm-translate")


def _system_prompt(target_language: str) -> str:
    """构造翻译系统提示词（返回单个字符串，不用隐式拼接避免 tuple bug）。

    内容：明确要求逐行独立翻译；碎片句（不成句的助词/名词/语气词）也要结合
    上下文给出自然中文并独立成行——这直接削弱 LLM 为求通顺而合并/拆分的倾向，
    是行数错位的主要诱发源。
    """
    return (
        "你是专业字幕翻译。将用户提供的日文字幕翻译为"
        + target_language
        + "。每行是一条独立字幕，必须逐行独立翻译。"
        + "有些行可能是不完整的日语碎片（单独的助词/名词/语气词），"
        + "请结合前后文语境给出它最自然的中文含义并独立成行。"
        + "输入有 N 行，输出就必须恰好 N 行中文、顺序保持一致。"
        + "绝对禁止把两行合并成一行，也禁止把一行拆成两行。"
        + "只返回译文，不要解释。"
    )


def _call_llm(
    api_base: str,
    api_key: str,
    model: str,
    system_prompt: str,
    user_content: str,
    request_timeout: float,
    log_prefix: str = "",
    **_: object,
) -> tuple[str, dict | None]:
    """发送一次 OpenAI 兼容的 chat.completions 请求，返回 (content, usage)。

    支持响应 choices[0].message.content 字段；enable_thinking=False 避免
    Qwen3 等模型的 reasoning_content 占满输出导致 content 为空/截断。

    usage 为响应体里的 usage 对象（含 prompt_tokens/completion_tokens/
    total_tokens），部分兼容接口不返回 usage 时为 None——调用方用其估算
    token 处理速度。log_prefix 为日志行前缀（如"第 2/5 批"），用于打印
    单批耗时与 token 速度。
    """
    started = time.monotonic()
    body = {
        "model": model,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_content},
        ],
        "enable_thinking": False,
        "max_tokens": 8192,
    }
    headers = {"Content-Type": "application/json"}
    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",
    )
    with urllib.request.urlopen(request, timeout=request_timeout) as response:
        payload = json.loads(response.read().decode("utf-8"))
    content = payload["choices"][0]["message"]["content"]
    usage = payload.get("usage")
    # 单批耗时与 token 速度日志：直观反映 LLM 处理速度（wall clock）。
    elapsed = time.monotonic() - started
    tokens = int(usage.get("total_tokens", 0)) if isinstance(usage, dict) else 0
    rate = tokens / elapsed if elapsed > 0 and tokens > 0 else 0.0
    logger.info(
        "LLM 响应 %s 耗时 %.1fs, tokens=%d (%.1f tok/s)",
        log_prefix, elapsed, tokens, rate,
    )
    return content, usage


def _repair_batch(batch: list[str], expected: int) -> list[str]:
    """把 LLM 返回的一个批次修整到与输入一致的行数（多合并、少补齐）。

    多行：末尾多出的行并入前一行（碎片本质同一句，时间轴落在该行窗口内）；
    少行：末尾补空串占位（宁缺勿错位，不挤占相邻字幕的时间轴）。
    """
    if len(batch) == expected:
        return batch
    if len(batch) > expected:
        merged = list(batch[:expected])
        merged[-1] = " ".join(batch[expected - 1 :])
        return merged
    # 少行补空串。
    return list(batch) + [""] * (expected - len(batch))


def translate_lines(lines: list[str], params: dict) -> list[str]:
    """分批调用 LLM 翻译纯文本行，返回顺序一致的译文列表。

    每批输入行数保持一致；若 LLM 返回行数不一致：多行合并、少行重试该批
    （最多 MAX_BATCH_RETRIES 次），仍不足则补空串占位。保证每条字幕都有
    译文且时间轴与原文逐条对齐，杜绝"内容对错时间"的错位。

    日志：每完成一批打印总进度（已完成行数/总行数、第几批/共几批、累计
    耗时与行处理速度），结束打印汇总（总耗时、累计 tokens 与 tok/s），
    便于评估 LLM 处理速度。
    """
    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", "600"))
    model = str(params.get("model") or os.getenv("LLM_MODEL", "Qwen/Qwen3.6-35B-A3B"))
    target_language = str(params.get("target_language", "zh-CN"))
    system_prompt = _system_prompt(target_language)

    total_lines = len(lines)
    total_batches = (total_lines + CHUNK_SIZE - 1) // CHUNK_SIZE if total_lines else 0
    if total_lines == 0:
        return []
    # 任务开始日志：总行数与总批数（批次 = CHUNK_SIZE 行，最后一批可能不足）。
    logger.info("翻译开始: %d 行, 分 %d 批", total_lines, total_batches)

    translated: list[str] = []
    total_tokens = 0
    all_started = time.monotonic()
    for batch_index in range(1, total_batches + 1):
        start = (batch_index - 1) * CHUNK_SIZE
        chunk = lines[start : start + CHUNK_SIZE]
        # 每批日志前缀（第几批/共几批），供单次 LLM 请求日志与批进度复用。
        log_prefix = f"第 {batch_index}/{total_batches} 批"
        batch_started = time.monotonic()
        batch_translated, batch_tokens = _translate_batch(
            chunk, api_base, api_key, model, system_prompt, request_timeout, log_prefix
        )
        translated.extend(batch_translated)
        total_tokens += batch_tokens
        # 批进度日志：已完成行数/总行数、当前批耗时、累计耗时与行处理速度。
        done = len(translated)
        elapsed_total = time.monotonic() - all_started
        logger.info(
            "翻译进度 %d/%d 行 (%s完成, 批耗时 %.1fs, 累计 %.1fs, %.1f 行/s)",
            done, total_lines, log_prefix,
            time.monotonic() - batch_started, elapsed_total,
            done / elapsed_total if elapsed_total > 0 else 0.0,
        )
    # 任务汇总日志：总耗时、累计 tokens 与 token/行处理速度。
    wall = time.monotonic() - all_started
    tok_rate = total_tokens / wall if wall > 0 and total_tokens > 0 else 0.0
    logger.info(
        "翻译完成: %d/%d 行, %d 批, 总耗时 %.1fs, 累计 tokens=%d (%.1f tok/s, %.1f 行/s)",
        len(translated), total_lines, total_batches, wall,
        total_tokens, tok_rate,
        len(translated) / wall if wall > 0 else 0.0,
    )
    return translated


def _translate_batch(
    chunk: list[str],
    api_base: str,
    api_key: str,
    model: str,
    system_prompt: str,
    request_timeout: float,
    log_prefix: str = "",
) -> tuple[list[str], int]:
    """翻译单个批次，返回 (与 chunk 等长译文, 本批 total_tokens)。

    行数不一致时多行合并、少行重试；每批调用前根据本批原文命中情况动态
    拼接专名/隐语规则（build_proper_noun_rule），注入到系统提示词，让 LLM
    正确处理片假名专名与成人语境隐语。"""
    # 本批命中的专名/隐语规则（无命中返回 None）。
    rule = build_proper_noun_rule(chunk)
    batch_system = system_prompt
    if rule:
        batch_system = system_prompt + "\n\n" + rule
    attempt = 0
    batch_tokens = 0
    while True:
        # content 为译文文本；usage 含本批 prompt/completion tokens（接口不
        # 返回时为 None），用于累计任务 token 总量与速度评估。
        content, usage = _call_llm(
            api_base,
            api_key,
            model,
            batch_system,
            "\n".join(chunk),
            request_timeout,
            log_prefix,
        )
        if isinstance(usage, dict):
            batch_tokens = int(usage.get("total_tokens", 0) or 0)
        # 保留所有行：先 rstrip 尾随换行避免多出末尾空行，再 splitlines 保留
        # 内容中的空串行（空行可能是合法的空字幕，过滤掉会误判行数）。
        batch = content.rstrip("\n").splitlines()
        if len(batch) == len(chunk):
            return batch, batch_tokens
        if len(batch) > len(chunk):
            # 多行：末尾多出的行合并到前一行，直接返回。
            return _repair_batch(batch, len(chunk)), batch_tokens
        # 少行：内容缺失，占位补空会丢语义，重试本批。
        attempt += 1
        if attempt >= MAX_BATCH_RETRIES:
            # 重试耗尽：补空串占位（宁缺勿错位），避免整条任务失败。
            return _repair_batch(batch, len(chunk)), batch_tokens


def invoke(request: InvokeRequest) -> InvokeResponse:
    """翻译 SRT 文件中的字幕文本，输出 cn.srt。"""
    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")

    # 标准 SRT 每 4 行一组：序号、时间轴、文本、空行；文本位于第 3 行。
    lines = srt_path.read_text(encoding="utf-8").splitlines()
    text_indices = list(range(2, len(lines), 4))
    source_lines = [lines[index] for index in text_indices]

    # 翻译：返回与 source_lines 严格等长的译文（多/少行已在批内修复）。
    translated_lines = translate_lines(source_lines, request.params)
    # 防御性兜底：确保长度一致（translate_lines 已保证，此处双保险）。
    translated_lines = translated_lines[: len(source_lines)]
    translated_lines += [""] * max(0, len(source_lines) - len(translated_lines))

    # 只替换文本行，序号、时间轴和空行保持不变。
    for index, text_index in enumerate(text_indices):
        lines[text_index] = translated_lines[index]

    # 长时寒暄幻觉词清洗：对展示时长超过阈值且含收尾/开场寒暄（晚安、感谢观看
    # 等）的条目，**连带时间戳整条删除**（剩余重编号），避免幻觉占位污染正片/ASS；
    # 短时（≤阈值）如剧情中真实互道'晚安'则保留，不误删。见
    # nodes/subtitle_cleanup.py。
    srt_body = "\n".join(lines) + "\n"
    srt_body = clean_srt_text(srt_body)

    output_dir = Path(request.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / "cn.srt"
    output_path.write_text(srt_body, encoding="utf-8")
    return InvokeResponse(status="completed", outputs={"cn_srt_uri": str(output_path)})