from __future__ import annotations

import argparse

from .infer import load_generator, write_wave


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="wfloat-tts")
    parser.add_argument("--model", "--checkpoint", dest="model", default="model.safetensors")
    parser.add_argument("--config", default="config.json")
    parser.add_argument("--text", required=True)
    parser.add_argument("--sid", type=int, default=0)
    parser.add_argument("--emotion", default="neutral")
    parser.add_argument("--intensity", type=float, default=0.5)
    parser.add_argument("--noise-scale", type=float, default=None)
    parser.add_argument("--length-scale", type=float, default=None)
    parser.add_argument("--noise-w", type=float, default=None)
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--output", required=True)
    return parser


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    generator = load_generator(
        checkpoint_path=args.model,
        config_path=args.config,
        device=args.device,
    )
    audio = generator.generate(
        text=args.text,
        sid=args.sid,
        emotion=args.emotion,
        intensity=args.intensity,
        noise_scale=args.noise_scale,
        length_scale=args.length_scale,
        noise_w=args.noise_w,
    )
    write_wave(args.output, audio.samples, audio.sample_rate)


if __name__ == "__main__":
    main()