#!/usr/bin/env python3
"""Export a pocket-tts voice state from a WAV sample into tts_voices.json.

Must run with <.venv-pockettts>/bin/python (torch + pocket-tts installed
there), never with the main .venv:

    .venv-pockettts/bin/python tools/pocket_tts_add_voice.py \\
        --wav /path/to/sample.wav --name "Marius"

Determines the next free index for --name by scanning tts_voices.json (an
entry counts as free if it's absent, or present but its state_path file
doesn't exist yet on disk), exports the voice state to
models/voices/<slug>.safetensors, and registers it in tts_voices.json via
src.jsonlock.
"""
import argparse
import re
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent
CATALOG_PATH = REPO_ROOT / "tts_voices.json"
VOICES_DIR = REPO_ROOT / "models" / "voices"

sys.path.insert(0, str(REPO_ROOT))

from src.credential_store import get_secret  # noqa: E402
from src.jsonlock import read_locked, write_locked  # noqa: E402


def _slugify(name: str) -> str:
    return re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-")


def _next_free_index(catalog: dict, name: str) -> int:
    n = 1
    while True:
        entry = catalog.get(f"{name} {n}")
        if entry is None:
            return n
        state_path = entry.get("state_path")
        if state_path and not (REPO_ROOT / state_path).exists():
            return n
        n += 1


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--wav", required=True, help="Path to source WAV sample")
    parser.add_argument("--name", required=True, help="Base voice name, e.g. 'Marius'")
    args = parser.parse_args()

    wav_path = Path(args.wav).expanduser().resolve()
    if not wav_path.exists():
        print(f"[pocket_tts_add_voice] ERROR: wav not found: {wav_path}", file=sys.stderr)
        return 1

    import os

    hf_token = get_secret("hf_token")
    if not hf_token:
        print(
            "[pocket_tts_add_voice] ERROR: hf_token lipsește din keyring (service echo-core).",
            file=sys.stderr,
        )
        return 1
    os.environ["HF_TOKEN"] = hf_token

    from pocket_tts import TTSModel, export_model_state

    catalog = read_locked(str(CATALOG_PATH))
    index = _next_free_index(catalog, args.name)
    full_name = f"{args.name} {index}"
    slug = _slugify(full_name)
    VOICES_DIR.mkdir(parents=True, exist_ok=True)
    export_path = VOICES_DIR / f"{slug}.safetensors"

    print("[pocket_tts_add_voice] Loading pocket-tts model...")
    model = TTSModel.load_model()
    print(f"[pocket_tts_add_voice] Exporting voice state from {wav_path} -> {export_path}")
    state = model.get_state_for_audio_prompt(audio_conditioning=str(wav_path), truncate=True)
    export_model_state(state, str(export_path))

    def _register(data: dict) -> dict:
        data[full_name] = {
            "engine": "pockettts",
            "state_path": f"models/voices/{slug}.safetensors",
            "owner": args.name,
        }
        return data

    write_locked(str(CATALOG_PATH), _register)
    print(f"[pocket_tts_add_voice] Registered voice '{full_name}' -> {export_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())