src/omnisub/profiles.py

"""Bước 3 — Hồ sơ giọng (VoiceProfile).

Đưa các clip audio của từng speaker cho Qwen3-Omni "nghe" → suy ra giới tính,
độ tuổi, cảm xúc nền, vai vế và gợi ý xưng hô tiếng Việt. Thay cho wav2vec2
age/gender + regex của bản cũ.
"""

from __future__ import annotations

import json
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Dict, List, Optional

from .backends.base import ChatMessage, LLMBackend


@dataclass
class VoiceProfile:
    speaker: str                 # "SPEAKER_01" (khớp pyannote)
    gender: str = "chưa rõ"      # nam | nữ | trẻ em | chưa rõ
    age_range: str = "chưa rõ"   # trẻ em | thiếu niên | thanh niên | trung niên | lớn tuổi
    emotion_baseline: str = ""   # vd: điềm đạm / nóng nảy / dịu dàng
    role_guess: str = ""         # vd: "phụ nữ trẻ, có vẻ là người yêu của SPEAKER_02"
    register_hint: str = ""      # gợi ý xưng hô VN
    evidence: str = ""           # model tự giải thích ngắn

    def to_prompt_line(self) -> str:
        bits = [f"{self.speaker}: {self.gender}, {self.age_range}"]
        if self.emotion_baseline:
            bits.append(f"giọng {self.emotion_baseline}")
        if self.role_guess:
            bits.append(self.role_guess)
        if self.register_hint:
            bits.append(f"xưng hô: {self.register_hint}")
        return " — ".join(bits)


_PROFILE_SYSTEM = (
    "Bạn là chuyên gia phân tích giọng nói cho phim Hoa ngữ. "
    "Nghe các clip của MỘT nhân vật rồi mô tả họ. "
    "Trả về DUY NHẤT một object JSON, không kèm giải thích ngoài JSON."
)


def _build_profile_prompt(speaker: str, n_clips: int) -> str:
    return (
        f"Đây là {n_clips} clip giọng của nhân vật {speaker} trong một phim Trung Quốc.\n"
        "Hãy nghe và suy ra hồ sơ nhân vật. Trả về JSON theo schema:\n"
        "{\n"
        '  "gender": "nam|nữ|trẻ em|chưa rõ",\n'
        '  "age_range": "trẻ em|thiếu niên|thanh niên|trung niên|lớn tuổi",\n'
        '  "emotion_baseline": "mô tả ngắn tông giọng nền",\n'
        '  "role_guess": "đoán vai vế/quan hệ nếu có manh mối",\n'
        '  "register_hint": "gợi ý cách xưng hô tiếng Việt phù hợp",\n'
        '  "evidence": "1 câu lý do dựa trên đặc điểm giọng"\n'
        "}"
    )


def build_voice_profile(
    backend: LLMBackend,
    speaker: str,
    clip_paths: List[str | Path],
    *,
    max_new_tokens: int = 512,
    **sampling,
) -> VoiceProfile:
    """Gọi Omni nghe các clip của 1 speaker → VoiceProfile."""
    audio = [str(p) for p in clip_paths]
    msg = ChatMessage(
        role="user",
        text=_build_profile_prompt(speaker, len(audio)),
        audio=audio,
    )
    try:
        data = backend.chat_json(
            [msg], system=_PROFILE_SYSTEM, max_new_tokens=max_new_tokens, **sampling
        )
    except Exception as e:  # model lỗi/parse lỗi → hồ sơ rỗng, không chặn pipeline
        return VoiceProfile(speaker=speaker, evidence=f"(không phân tích được: {e})")

    if not isinstance(data, dict):
        return VoiceProfile(speaker=speaker)
    return VoiceProfile(
        speaker=speaker,
        gender=str(data.get("gender", "chưa rõ")),
        age_range=str(data.get("age_range", "chưa rõ")),
        emotion_baseline=str(data.get("emotion_baseline", "")),
        role_guess=str(data.get("role_guess", "")),
        register_hint=str(data.get("register_hint", "")),
        evidence=str(data.get("evidence", "")),
    )


def build_all_profiles(
    backend: LLMBackend,
    speaker_clips: Dict[str, List[str | Path]],
    *,
    max_new_tokens: int = 512,
    **sampling,
) -> Dict[str, VoiceProfile]:
    profiles: Dict[str, VoiceProfile] = {}
    for speaker, clips in speaker_clips.items():
        if not clips:
            profiles[speaker] = VoiceProfile(speaker=speaker)
            continue
        profiles[speaker] = build_voice_profile(
            backend, speaker, clips, max_new_tokens=max_new_tokens, **sampling
        )
    return profiles


def save_profiles(profiles: Dict[str, VoiceProfile], path: str | Path) -> None:
    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(
        json.dumps({k: asdict(v) for k, v in profiles.items()}, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )


def load_profiles(path: str | Path) -> Dict[str, VoiceProfile]:
    data = json.loads(Path(path).read_text(encoding="utf-8"))
    return {k: VoiceProfile(**v) for k, v in data.items()}


def profiles_prompt_block(profiles: Dict[str, VoiceProfile]) -> str:
    """Khối text mô tả hồ sơ giọng để nhét vào prompt dịch."""
    if not profiles:
        return ""
    lines = [p.to_prompt_line() for p in profiles.values()]
    return "Hồ sơ nhân vật (suy từ giọng nói):\n" + "\n".join(f"- {ln}" for ln in lines)