"""Bước 4 — Dịch đa phương thức theo cảnh.

Dựng prompt gồm: frame (trước text) + cues ZH + hồ sơ giọng + glossary (text)
+ audio cảnh (sau text). Gọi Omni, parse JSON từng cue, ghi `cue.translation`
và `cue.note`. KHÔNG dùng regex xưng hô.
"""

from __future__ import annotations

from typing import Dict, List, Optional

from .backends.base import ChatMessage, LLMBackend
from .glossary import Glossary
from .profiles import VoiceProfile, profiles_prompt_block
from .scene_context import SceneContext
from .scenes import Scene
from .srt import subtitle_char_budget


TRANSLATE_SYSTEM = (
    "Bạn là dịch giả phụ đề phim Hoa ngữ sang tiếng Việt, văn nói tự nhiên, súc tích. "
    "Bạn ĐƯỢC nghe audio và xem khung hình của cảnh để chọn xưng hô đúng vai vế "
    "(anh/em, chị/em, mẹ/con, ngài/ta, tỷ/muội...). "
    "Giữ nhất quán tên riêng và cách xưng hô theo glossary đã cho. "
    "Trả về DUY NHẤT một object JSON đúng schema, không thêm chữ nào ngoài JSON."
)


def _cue_block(scene: Scene, profiles: Dict[str, VoiceProfile]) -> str:
    lines = []
    for c in scene.cues:
        spk = c.speaker or "?"
        budget = subtitle_char_budget(c.duration)
        lines.append(f'  {{ "id": {c.index}, "speaker": "{spk}", "zh": {c.text!r}, "max_chars": {budget} }}')
    return "[\n" + ",\n".join(lines) + "\n]"


def build_scene_messages(
    scene: Scene,
    context: Optional[SceneContext],
    profiles: Dict[str, VoiceProfile],
    glossary: Glossary,
    *,
    source_lang: str = "Chinese",
    target_lang: str = "Vietnamese",
) -> List[ChatMessage]:
    """Dựng message đa phương thức cho một cảnh (ảnh trước, audio sau text)."""
    images = [str(p) for p in (context.frame_paths if context else [])]
    audio = [str(context.audio_path)] if (context and context.audio_path) else []

    prompt_parts: List[str] = [
        f"Dịch các câu thoại sau từ {source_lang} sang {target_lang}.",
        f"Cảnh #{scene.scene_id}, thời lượng ~{scene.duration:.1f}s.",
    ]
    prof_block = profiles_prompt_block(profiles)
    if prof_block:
        prompt_parts.append(prof_block)
    gloss_block = glossary.prompt_block()
    if gloss_block:
        prompt_parts.append(gloss_block)

    prompt_parts.append("Các câu cần dịch (giữ đúng id):\n" + _cue_block(scene, profiles))
    prompt_parts.append(
        "Yêu cầu:\n"
        "- Dịch tự nhiên, KHÔNG vượt quá max_chars mỗi câu (rút gọn nếu cần).\n"
        "- Chọn xưng hô dựa vào giọng nói + hình ảnh + ngữ cảnh.\n"
        "- Đề xuất tên riêng và cặp xưng hô để khóa nhất quán cho cảnh sau.\n"
        "Trả về JSON:\n"
        "{\n"
        '  "cues": [{ "id": <int>, "vi": "<bản dịch>", "note": "<ghi chú xưng hô ngắn>" }],\n'
        '  "names": [{ "source": "<tên ZH>", "target": "<tên VN>" }],\n'
        '  "address": [{ "from": "SPEAKER_X", "to": "SPEAKER_Y", "address": "anh-em" }]\n'
        "}"
    )
    if audio:
        prompt_parts.append("(Audio cảnh được đính kèm phía dưới để bạn nghe.)")

    msg = ChatMessage(
        role="user",
        text="\n\n".join(prompt_parts),
        images=images,
        audio=audio,
    )
    return [msg]


def apply_scene_result(scene: Scene, result: dict) -> None:
    """Ghi bản dịch + note vào từng cue theo id."""
    by_id = {c.index: c for c in scene.cues}
    for item in (result.get("cues", []) if isinstance(result, dict) else []) or []:
        cid = item.get("id")
        if cid in by_id:
            vi = item.get("vi") or item.get("translation")
            if vi:
                by_id[cid].translation = str(vi).strip()
            note = item.get("note")
            if note:
                by_id[cid].note = str(note).strip()


def translate_scene(
    backend: LLMBackend,
    scene: Scene,
    context: Optional[SceneContext],
    profiles: Dict[str, VoiceProfile],
    glossary: Glossary,
    *,
    source_lang: str = "Chinese",
    target_lang: str = "Vietnamese",
    max_new_tokens: int = 1024,
    **sampling,
) -> dict:
    """Dịch một cảnh: gọi Omni, ghi kết quả vào cue, cập nhật glossary."""
    messages = build_scene_messages(
        scene, context, profiles, glossary,
        source_lang=source_lang, target_lang=target_lang,
    )
    try:
        result = backend.chat_json(
            messages, system=TRANSLATE_SYSTEM, max_new_tokens=max_new_tokens, **sampling
        )
    except Exception as e:
        # fallback: giữ nguyên text gốc, đánh dấu lỗi
        for c in scene.cues:
            c.note = f"(lỗi dịch: {e})"
        return {}

    if not isinstance(result, dict):
        result = {}
    apply_scene_result(scene, result)
    glossary.update_from_scene_result(result)
    return result