"""LLM 翻译节点入口。

通过标准节点 HTTP 服务接收 SRT，提取纯文本行分批调用 LLM，再把译文回填到
原 SRT 结构并输出 cn.srt。
"""

from __future__ import annotations

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

from wov_sdk.models import InvokeRequest, InvokeResponse, NodeManifest
from wov_sdk.server import run_node

# 单次 LLM 请求携带的字幕行数；过大会超出模型上下文，过小则请求次数过多。
CHUNK_SIZE = 20


def translate_lines(lines: list[str], params: dict) -> list[str]:
    """分批调用 LLM 翻译纯文本行，返回顺序一致的译文列表。"""
    # 接口地址、Key 和模型均可通过环境变量配置，默认指向内网兼容接口。
    api_base = os.getenv(
        "LLM_API_BASE",
        "http://192.168.123.70:8080/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", "default"))
    target_language = str(params.get("target_language", "zh-CN"))
    # 系统提示词约束模型只输出译文，保证行数和顺序可回填。
    system_prompt = (
        "你是专业字幕翻译。将用户提供的日文字幕翻译为"
        f"{target_language}。只返回译文，保持行数和顺序，不要添加解释。"
    )
    translated: list[str] = []
    # 按 CHUNK_SIZE 分批发送，避免单次请求超过模型上下文限制。
    for start in range(0, len(lines), CHUNK_SIZE):
        chunk = lines[start : start + CHUNK_SIZE]
        body = {
            "model": model,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": "\n".join(chunk)},
            ],
        }
        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",
        )
        with urllib.request.urlopen(request, timeout=request_timeout) as response:
            payload = json.loads(response.read().decode("utf-8"))
        # 兼容 OpenAI Chat Completions 响应格式，取第一条消息内容。
        content = payload["choices"][0]["message"]["content"]
        # 忽略空行，保证译文列表与输入行一一对应。
        translated.extend(
            [line.strip() for line in content.splitlines() if line.strip()]
        )
    return translated


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]
    translated_lines = translate_lines(source_lines, request.params)
    # 防止模型返回行数偏差：多出的截断，缺少的用空串补齐。
    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]

    output_dir = Path(request.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / "cn.srt"
    # 末尾补一个换行，让文件满足常见文本工具习惯。
    output_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return InvokeResponse(status="completed", outputs={"cn_srt_uri": str(output_path)})


def main() -> None:
    """加载节点清单并以本模块的 invoke 处理器启动服务。"""
    manifest_path = Path(__file__).resolve().parent.parent / "node.manifest.json"
    with open(manifest_path, "r", encoding="utf-8") as f:
        manifest = NodeManifest.from_dict(json.load(f))
    run_node(manifest, invoke)


if __name__ == "__main__":
    main()