import sys
import time
from pathlib import Path
import sherpa_onnx, soundfile as sf
from core import tts_room as ap
from config.tts import (
    TOKENS, ENCODER, DECODER, DATA_DIR, LEXICON, VOCODER,
    REFERENCE_AUDIO, REFERENCE_TEXT, NUM_THREADS, PROVIDER, NUM_STEPS, OUTPUT_FILE,
)

def tts_run(args):

    ref_audio = Path(args.ref_audio) if args.ref_audio is not None else Path(REFERENCE_AUDIO)
    ref_text  = args.ref_text       if args.ref_text  is not None else REFERENCE_TEXT

    if not ref_audio.is_file():
        print(f"Reference audio not found: {ref_audio}", file=sys.stderr)
        return

    print("Loading model...")
    tts_config = sherpa_onnx.OfflineTtsConfig(
        model=sherpa_onnx.OfflineTtsModelConfig(
            zipvoice=sherpa_onnx.OfflineTtsZipvoiceModelConfig(
                tokens   = str(TOKENS),
                encoder  = str(ENCODER),
                decoder  = str(DECODER),
                data_dir = str(DATA_DIR),
                lexicon  = str(LEXICON),
                vocoder  = str(VOCODER),
            ),
            debug       = False,
            num_threads = NUM_THREADS,
            provider    = PROVIDER,
        )
    )
    if not tts_config.validate():
        raise ValueError("Invalid TTS config. Please read the previous error messages.")

    tts = sherpa_onnx.OfflineTts(tts_config)
    print("Model ready!")

    reference_audio, sr = sf.read(str(ref_audio), dtype="float32")
    if reference_audio.ndim > 1:
        reference_audio = reference_audio[:, 0]

    gen_config = sherpa_onnx.GenerationConfig()
    gen_config.reference_audio       = reference_audio
    gen_config.reference_sample_rate = sr
    gen_config.reference_text        = ref_text
    gen_config.num_steps             = NUM_STEPS

    output_file = args.output if args.output is not None else OUTPUT_FILE

    print("Generating...")
    start = time.time()
    audio = tts.generate(args.text, gen_config)
    elapsed = time.time() - start

    if len(audio.samples) == 0:
        print("Error in generating audio. Please read the previous error messages.", file=sys.stderr)
        return

    sf.write(output_file, audio.samples, samplerate=audio.sample_rate, subtype="PCM_16")

    duration = len(audio.samples) / audio.sample_rate
    print()
    print(f"Output       : {output_file}")
    print(f"Text         : {args.text}")
    print(f"Duration     : {duration:.3f} s")
    print(f"Elapsed      : {elapsed:.3f} s")
    print(f"RTF          : {elapsed:.3f} / {duration:.3f} = {elapsed / duration:.3f}")
    print()

if __name__ == "__main__":
    tts_run( ap.parser.parse_args() )