gart/admin/api/transcription.py

"""Transcrição e locutores.

Extraído de models_api.py — a tabela de comandos segue lá.
"""

from __future__ import annotations

import json
from pathlib import Path

from fcpxml.diarize import (
    assign_speakers,
    build_speakers,
    diarization_capability,
    diarize,
)
from fcpxml.media_intel import media_src_to_path
from fcpxml.model_manager import (
    is_model_downloaded,
    load_hf_token,
    load_num_speakers,
    load_selected_model,
    load_transcript_language,
    save_hf_token,
    save_num_speakers,
)
from fcpxml.parser import parse_fcpxml
from fcpxml.transcribe import transcribe

from . import shared
from .shared import (
    _load_cached_transcript,
    _save_json_atomic,
    _transcript_json_path,
)


def cmd_transcribe(args: dict) -> int:
    proj_path = str(args.get("path", ""))
    output_dir = str(args.get("output_dir", "")).strip()
    # Honra o modelo selecionado no programa quando nenhum é passado.
    model = str(args.get("model", "") or load_selected_model() or "")
    language = args.get("language")
    if language is None:
        language = load_transcript_language()
    if language == "auto":
        language = None
    if not proj_path:
        shared.emit({"type": "error", "message": "Nenhum projeto selecionado."})
        return 1
    if not output_dir:
        shared.emit({"type": "error", "message": "Selecione a pasta do projeto antes de transcrever."})
        return 1
    if not model or not is_model_downloaded(model):
        shared.emit(
            {
                "type": "error",
                "message": "Nenhum modelo de transcrição instalado. Baixe e selecione um modelo na aba Modelos.",
            }
        )
        return 1

    token = str(args.get("hf_token") or load_hf_token() or "")
    if args.get("num_speakers") is not None:
        num_speakers = str(args.get("num_speakers"))
    else:
        num_speakers = load_num_speakers()

    # Load project.
    try:
        proj = parse_fcpxml(proj_path)
    except Exception as exc:
        shared.emit({"type": "error", "message": f"Erro ao ler o projeto: {exc}"})
        return 1
    tl = proj.primary_timeline or (proj.timelines[0] if proj.timelines else None)
    media_paths: list[str] = []
    if tl is not None:
        for clip in getattr(tl, "clips", []):
            mp = media_src_to_path(clip.media_path or "")
            if mp and Path(mp).is_file() and mp not in media_paths:
                media_paths.append(mp)
    if not media_paths:
        shared.emit({"type": "error", "message": "Nenhum arquivo de mídia acessível encontrado."})
        return 1

    total = len(media_paths)
    results: list[dict] = []
    for i, mp in enumerate(media_paths, 1):
        stage = f"Transcrevendo {Path(mp).name} ({i}/{total})…"
        shared.emit({"type": "progress", "fraction": (i - 1) / total, "stage": stage})
        json_path = _transcript_json_path(mp, output_dir)
        cached = _load_cached_transcript(json_path)
        if cached is not None:
            shared.emit({"type": "progress", "fraction": i / total, "stage": stage})
            results.append(_result_row(mp, cached))
            continue

        def _on_progress(file_fraction: float, _i: int = i, _stage: str = stage) -> None:
            # Blend this file's own progress into the overall fraction so a
            # single-media project doesn't jump straight to 100% before the
            # actual (slow) decoding work has even started.
            overall = (_i - 1 + file_fraction) / total
            shared.emit({"type": "progress", "fraction": overall, "stage": _stage})

        data = transcribe(mp, model_size=model, language=language, progress_cb=_on_progress)
        if data is None:
            shared.emit({"type": "error", "message": f"Não foi possível transcrever: {Path(mp).name}"})
            return 1

        # Diarização opcional (necessita token HF): assina speaker por segmento/palavra.
        if token:
            tracks = diarize(mp, token, num_speakers)
            segments, words = assign_speakers(
                data.get("segments", []), data.get("words", []), tracks
            )
            data = {**data, "segments": segments, "words": words}
        data["speakers"] = build_speakers(data.get("segments", []))

        payload = {
            "schema_version": "1.0",
            "source": Path(mp).name,
            "model": model,
            **data,
        }
        try:
            _save_json_atomic(json_path, payload)
        except (OSError, RuntimeError, ValueError) as exc:
            shared.emit({"type": "error", "message": f"Não foi possível salvar o JSON: {exc}"})
            return 1
        results.append(_result_row(mp, data))

    shared.emit({"type": "result", "transcripts": results})
    return 0


def cmd_rename_speakers(args: dict) -> int:
    """Apply real names to speakers already saved in a transcript JSON."""
    json_path = Path(str(args.get("path", "")))
    names = args.get("speakers") or {}
    if not json_path.is_file():
        shared.emit({"type": "error", "message": "Transcrição não encontrada."})
        return 1
    try:
        data = json.loads(json_path.read_text(encoding="utf-8"))
    except (OSError, ValueError) as exc:
        shared.emit({"type": "error", "message": f"Não foi possível ler o JSON: {exc}"})
        return 1
    mapping = {str(sid): str(name).strip() for sid, name in (names or {}).items()}
    for sp in data.get("speakers", []):
        sid = str(sp.get("id", ""))
        if mapping.get(sid):
            sp["name"] = mapping[sid]
    try:
        _save_json_atomic(json_path, data)
    except (OSError, RuntimeError, ValueError) as exc:
        shared.emit({"type": "error", "message": f"Não foi possível salvar: {exc}"})
        return 1
    shared.emit({"ok": True, "speakers": data.get("speakers", [])})
    return 0

def cmd_set_diarization(args: dict) -> int:
    """Persist the HuggingFace token and expected speaker count for diarization."""
    token = args.get("token")
    num = args.get("num_speakers")
    if token is not None:
        save_hf_token(str(token))
    if num is not None:
        save_num_speakers(str(num))
    ok, msg = diarization_capability(load_hf_token())
    shared.emit({"ok": True, "diarization": ok, "diarization_message": msg, "num_speakers": load_num_speakers()})
    return 0

def _result_row(mp: str, data: dict) -> dict:
    words = data.get("words", [])
    preview = (data.get("text", "") or "")[:160]
    speakers = data.get("speakers") or []
    return {
        "media": Path(mp).name,
        "language": data.get("language", "?"),
        "words": len(words),
        "duration": float(data.get("duration", 0.0)),
        "preview": preview,
        "saved": str(_transcript_json_path(mp)),
        "speakers": [s.get("name", s.get("id", "")) for s in speakers],
    }