#!/usr/bin/env python3
import argparse
import sys
from pathlib import Path

import torch
import torchaudio

MODEL_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(MODEL_DIR))
sys.path.insert(0, str(MODEL_DIR / "third_party" / "Matcha-TTS"))

from accent_config import ACCENT_INSTRUCTIONS  # noqa: E402
from cosyvoice.cli.cosyvoice import AutoModel  # noqa: E402


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--text", required=True)
    parser.add_argument("--accent", choices=ACCENT_INSTRUCTIONS, default="singapore")
    parser.add_argument("--prompt_wav", default=str(MODEL_DIR / "zero_shot_prompt.wav"))
    parser.add_argument("--output", default="output.wav")
    parser.add_argument("--speed", type=float, default=1.0)
    parser.add_argument("--fp16", action=argparse.BooleanOptionalAction, default=torch.cuda.is_available())
    args = parser.parse_args()

    model = AutoModel(model_dir=str(MODEL_DIR), fp16=args.fp16, load_vllm=False)
    chunks = [item["tts_speech"] for item in model.inference_instruct2(
        args.text,
        ACCENT_INSTRUCTIONS[args.accent],
        args.prompt_wav,
        stream=False,
        speed=args.speed,
    )]
    speech = torch.cat(chunks, dim=1)
    torchaudio.save(args.output, speech.cpu(), model.sample_rate)
    print(args.output)


if __name__ == "__main__":
    main()