#!/usr/bin/env python3
"""Interactively generate and choose a dots.tts narrator reference."""

import argparse
import os
from pathlib import Path
import shutil
import subprocess
import sys
import time


DEFAULT_TEXT = "The story continues as our hero steps forward into the unknown."
DEFAULT_MODEL = "rednote-hilab/dots.tts-base"
DEFAULT_VOICE_DIR = Path("~/.cache/manga-tts").expanduser()


def load_model(model_name: str):
    # Match worker_tts.py's workaround for the tokenizer bundled with dots.tts.
    import transformers

    original = transformers.AutoTokenizer.from_pretrained.__func__
    transformers.AutoTokenizer.from_pretrained = classmethod(
        lambda cls, *args, **kwargs: original(
            cls, *args, **{"fix_mistral_regex": True, **kwargs}
        )
    )
    # Match worker_tts.py's torch/torchaudio minor-mismatch bypass too.
    import torch
    from importlib import metadata

    real_version = metadata.version
    metadata.version = (
        lambda name: torch.__version__ if name == "torchaudio" else real_version(name)
    )
    try:
        from dots_tts.runtime import DotsTtsRuntime
    finally:
        metadata.version = real_version

    return DotsTtsRuntime.from_pretrained(model_name, precision="bfloat16")


def write_wav(result, path: Path) -> None:
    import numpy as np
    import soundfile as sf

    audio = result["audio"]
    if hasattr(audio, "detach"):
        audio = audio.detach().cpu().numpy()
    sf.write(path, np.asarray(audio, dtype="float32").squeeze(),
             result["sample_rate"], subtype="PCM_16")


def play(path: Path) -> None:
    players = (
        ("ffplay", "-nodisp", "-autoexit", "-loglevel", "error"),
        ("aplay", "-q"),
        ("paplay",),
    )
    for command in players:
        if shutil.which(command[0]):
            subprocess.run([*command, str(path)], check=False)
            return
    print(f"No ffplay, aplay, or paplay found; listen manually: {path}")


def choice() -> str:
    while True:
        answer = input("[p]ick / p[a]rk / [s]kip: ").strip().lower()
        aliases = {"p": "pick", "pick": "pick", "a": "park", "park": "park",
                   "s": "skip", "skip": "skip"}
        if answer in aliases:
            return aliases[answer]
        print("Enter p, a, or s.")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--text", default=os.environ.get("VOICE_REF_TEXT", DEFAULT_TEXT),
                        help="sample text to speak")
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--voice-dir", type=Path, default=DEFAULT_VOICE_DIR)
    args = parser.parse_args()

    voice_dir = args.voice_dir.expanduser()
    parked_dir = voice_dir / "parked"
    voice_dir.mkdir(parents=True, exist_ok=True)
    parked_dir.mkdir(parents=True, exist_ok=True)
    candidate = voice_dir / "voice_candidate.wav"
    selected = voice_dir / "narrator_ref.wav"

    print(f"Loading {args.model} ...")
    model = load_model(args.model)
    print(f"Reference text: {args.text!r}")

    try:
        while True:
            print("\nGenerating a new voice ...")
            write_wav(model.generate(text=args.text), candidate)
            play(candidate)

            action = choice()
            if action == "pick":
                os.replace(candidate, selected)
                (voice_dir / "narrator_ref.txt").write_text(args.text + "\n")
                print(f"Selected: {selected}")
                return 0
            if action == "park":
                stamp = time.strftime("%Y%m%d-%H%M%S")
                parked = parked_dir / f"voice-{stamp}.wav"
                suffix = 2
                while parked.exists():
                    parked = parked_dir / f"voice-{stamp}-{suffix}.wav"
                    suffix += 1
                os.replace(candidate, parked)
                parked.with_suffix(".txt").write_text(args.text + "\n")
                print(f"Parked: {parked}")
            else:
                candidate.unlink(missing_ok=True)
                print("Skipped.")
    except (KeyboardInterrupt, EOFError):
        candidate.unlink(missing_ok=True)
        print("\nStopped without changing the selected voice.")
        return 130


if __name__ == "__main__":
    sys.exit(main())