"""
maven voice finder
generates single samples via VoiceDesign until you find one you like.
saves the winner as maven_reference.wav + its text as maven_reference.txt
for use with Base model cloning.

requires: pip install qwen-tts soundfile torch
also needs: sudo apt install ffmpeg  (for playback)
"""

import sys
import soundfile as sf
import subprocess
from pathlib import Path

import torch
from qwen_tts import Qwen3TTSModel

# ── config ────────────────────────────────────────────────────────────────────

MODEL     = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
REF_DIR = Path("ref")
TEXT_FILE = "maven_reference.txt"
AUDIO_FILE = "maven_reference.wav"

# neutral sentence — long enough to capture voice characteristics
SAMPLE_TEXT = (
    "Ну, если честно, я не уверена, что это сработает именно так. "
    "Но давай попробуем - хуже уже точно не будет."
    "Хотя, возможно, стоит подумать о другом подходе."
    "Ошибки нынче стоят дорого. "
)

VOICE_DESIGN = (
    "Female voice, mid-20s. Natural and grounded, but lively. "
    "Fast, slightly uneven pacing — thoughts often move ahead of speech. "
    "Light, reactive expressiveness: small spikes of interest, brief amused or curious tones, "
    "then returning to a calm baseline. "
    "Occasional micro-pauses, restarts, or cut-off phrases, like thinking in real time. "
    "Tone is clear and fairly light, but not high-pitched or cutesy. "
    "Not performative — feels spontaneous rather than acted. "
    "Subtle unpredictability in rhythm and delivery. "
    "No anime-style, no exaggerated эмоции, no theatrical voice acting."
)

# ── playback ──────────────────────────────────────────────────────────────────

def play(path: Path):
    """play wav via ffplay (ffmpeg). silent if not available."""
    try:
        subprocess.run(
            ["ffplay", "-nodisp", "-autoexit", "-loglevel", "quiet", str(path)],
            check=True
        )
    except FileNotFoundError:
        print("  [!] ffplay not found — open the wav manually")
    except subprocess.CalledProcessError:
        pass

# ── main ──────────────────────────────────────────────────────────────────────

def main():
    print(f"[*] loading model: {MODEL}")
    model = Qwen3TTSModel.from_pretrained(
        MODEL,
        device_map="cuda:0",
        dtype=torch.float32,
    )

    attempt = 0
    
    current_index = max([int(dir.stem) for dir in REF_DIR.iterdir() if dir.is_dir()])
    next_index = current_index + 1

    print(f"\ncurrent index is: {current_index}")
    print(f"\nnext will be: {next_index}")

    while True:
        attempt += 1
        print(f"\n[{attempt}] generating sample...")

        wavs, sr = model.generate_voice_design(
            text=SAMPLE_TEXT,
            instruct=VOICE_DESIGN,
            language="Russian",
        )
        audio = wavs[0] if isinstance(wavs, list) else wavs

        tmp = Path(f"candidate_{attempt:03d}.wav")
        sf.write(str(tmp), audio, sr)
        print(f"    saved → {tmp}")

        play(tmp)

        answer = input("    keep this voice? [y/n/q] ").strip().lower()

        if answer == "y":
            new_path = REF_DIR / str(next_index) / AUDIO_FILE
            if not new_path.exists():
                new_path.parent.mkdir()
            tmp.rename(new_path)
            text_path = REF_DIR / str(next_index) / TEXT_FILE
            text_path.write_text(SAMPLE_TEXT, encoding="utf-8")
            print(f"\n[✓] saved as {new_path}")
            print(f"[✓] reference text → {text_path}")
            print("\nnext step: update generate_dataset.py to use Base model + this reference.")
            tmp.unlink(missing_ok=True)
            break
        elif answer == "q":
            print("quitting.")
            tmp.unlink(missing_ok=True)
            sys.exit(0)
        else:
            print("    regenerating...")
            tmp.unlink(missing_ok=True)


if __name__ == "__main__":
    main()