"""Bước 1 (tùy chọn) — OCR sửa phụ đề cháy hình. MẶC ĐỊNH TẮT.

Khi `correct_ocr: true`, đưa frame + text OCR cho Omni để sửa lỗi nhận dạng
(ký tự sai, dính dòng). Nguồn SRT hiện đã sạch nên mặc định bỏ qua bước này.
"""

from __future__ import annotations

from typing import List, Optional

from .backends.base import ChatMessage, LLMBackend
from .scene_context import SceneContext
from .scenes import Scene


_CORRECT_SYSTEM = (
    "Bạn là công cụ sửa lỗi OCR phụ đề tiếng Trung cháy trên khung hình. "
    "So khớp văn bản với khung hình và sửa ký tự sai. "
    "Trả về JSON {\"cues\":[{\"id\":<int>,\"zh\":\"<text đã sửa>\"}]}, không thêm gì khác."
)


def correct_scene(
    backend: LLMBackend,
    scene: Scene,
    context: Optional[SceneContext],
    *,
    max_new_tokens: int = 1024,
    **sampling,
) -> None:
    """Sửa OCR text gốc của các cue trong cảnh (ghi đè `cue.text`)."""
    images = [str(p) for p in (context.frame_paths if context else [])]
    cue_lines = "\n".join(f'  {{ "id": {c.index}, "zh": {c.text!r} }}' for c in scene.cues)
    text = (
        "Dưới đây là phụ đề OCR và khung hình tương ứng. Sửa lỗi ký tự nếu có.\n"
        f"[\n{cue_lines}\n]"
    )
    msg = ChatMessage(role="user", text=text, images=images)
    try:
        result = backend.chat_json(
            [msg], system=_CORRECT_SYSTEM, max_new_tokens=max_new_tokens, **sampling
        )
    except Exception:
        return
    if not isinstance(result, dict):
        return
    by_id = {c.index: c for c in scene.cues}
    for item in result.get("cues", []) or []:
        cid = item.get("id")
        zh = item.get("zh")
        if cid in by_id and zh:
            by_id[cid].text = str(zh).strip()