"""Enrôlement vocal OmniVoice : WAV de référence → clone-prompt `.ovsp`.

⚠️ Sous l'env encodeur `ov_venv` (torch + omnivoice + transformers + librosa/soundfile),
invoqué par le worker subprocess (pas en-process API py3.14).
Réimplémente /opt/Kazeia-engine/dist/omnivoice/scripts/make_ovsp.py.
Format .ovsp : int32 'OVSP' | int32 T_ref | int32 text_len | int32 tokens[8*T_ref] | char ref_text[].
"""
from __future__ import annotations
import logging, os, struct, sys
import numpy as np

for _noisy in ("numba", "matplotlib", "transformers", "omnivoice", "huggingface_hub"):
    logging.getLogger(_noisy).setLevel(logging.WARNING)

OV_MODEL = os.environ.get("OV_MODEL", "k2-fsa/OmniVoice")
OVSP_MAGIC = 0x4F565350
# Cap ~16 s : T_ref ~390 -> bucket 640 (rapide). Au-delà = bucket 768/1024 (lent) +
# risque de saturer le plus grand bucket. prepare_wav partagé avec la transcription.
OV_TARGET_S = float(os.environ.get("OV_TARGET_S", "16.0"))
_PREP_SR = 24000
_HARD_MAX_S = 20.0
_MIN_WARN_S = 6.0
_model_singleton = None

class EnrollmentError(RuntimeError):
    pass

def prepare_wav(in_path: str, out_path: str, target_s: float = OV_TARGET_S) -> dict:
    import librosa, soundfile as sf
    y, _ = librosa.load(in_path, sr=_PREP_SR, mono=True)
    orig_s = len(y) / _PREP_SR
    yt, _ = librosa.effects.trim(y, top_db=30)
    if yt.size == 0:
        yt = y
    yt = yt[: int(min(target_s, _HARD_MAX_S) * _PREP_SR)]
    final_s = len(yt) / _PREP_SR
    sf.write(out_path, yt, _PREP_SR, subtype="PCM_16")
    return {"orig_s": round(orig_s, 1), "final_s": round(final_s, 1), "too_short": final_s < _MIN_WARN_S}

def load_model():
    global _model_singleton
    if _model_singleton is not None:
        return _model_singleton
    try:
        import torch
        from omnivoice.models.omnivoice import OmniVoice
    except Exception as e:
        raise EnrollmentError(f"import OmniVoice échoué ({e}). Exécuter sous ov_venv.") from e
    _model_singleton = OmniVoice.from_pretrained(OV_MODEL, torch_dtype=torch.float32).to("cpu").eval()
    return _model_singleton

def _clone_prompt(prep_wav: str, ref_text):
    import torch
    m = load_model()
    with torch.no_grad():
        return m.create_voice_clone_prompt(ref_audio=prep_wav, ref_text=ref_text or None)

def transcribe_ref(wav_path: str) -> dict:
    """ASR du segment préparé (16 s) via OmniVoice, pour relecture opérateur."""
    import tempfile
    if not os.path.isfile(wav_path):
        raise EnrollmentError(f"WAV introuvable: {wav_path}")
    with tempfile.TemporaryDirectory() as td:
        prep = os.path.join(td, "prep.wav")
        prep_info = prepare_wav(wav_path, prep)
        vcp = _clone_prompt(prep, None)
    return {"text": (getattr(vcp, "ref_text", "") or "").strip(),
            "language": None, "model": "omnivoice-asr", "prep": prep_info}

def enroll(wav_path: str, text: str, out_path: str) -> dict:
    import tempfile
    if not os.path.isfile(wav_path):
        raise EnrollmentError(f"WAV introuvable: {wav_path}")
    text = (text or "").strip()
    with tempfile.TemporaryDirectory() as td:
        prep = os.path.join(td, "prep.wav")
        prep_info = prepare_wav(wav_path, prep)
        vcp = _clone_prompt(prep, text if text else None)
    toks = vcp.ref_audio_tokens.cpu().numpy().astype(np.int32)   # [8, T_ref]
    C, T = toks.shape
    if C != 8:
        raise EnrollmentError(f"attendu 8 codebooks, eu {C}")
    ref_text = (getattr(vcp, "ref_text", "") or text or "").strip()
    txt = ref_text.encode("utf-8")
    os.makedirs(os.path.dirname(os.path.abspath(out_path)), exist_ok=True)
    with open(out_path, "wb") as f:
        f.write(struct.pack("<i", OVSP_MAGIC))
        f.write(struct.pack("<i", T))
        f.write(struct.pack("<i", len(txt)))
        f.write(toks.tobytes())
        f.write(txt)
    return {"out_path": out_path, "bytes": os.path.getsize(out_path), "t_ref": int(T),
            "ref_seconds": round(T / 25.0, 1), "text_bytes": len(txt),
            "ref_text": ref_text, "prep": prep_info}

if __name__ == "__main__":
    import json
    ref, ref_txt, out = sys.argv[1], sys.argv[2], sys.argv[3]
    text = "" if ref_txt == "-" else (open(ref_txt).read() if ref_txt.endswith(".txt") and os.path.isfile(ref_txt) else ref_txt)
    print(json.dumps(enroll(ref, text, out), ensure_ascii=False))