import os, sys, time
sys.path.insert(0, r'G:\supertonic\py')
import soundfile as sf
import onnxruntime as ort
from helper import load_cfgs, load_text_processor, load_onnx_all, load_voice_style, TextToSpeech, timer

print("ORT version:", ort.__version__)
print("Available providers:", ort.get_available_providers())
onnx_dir = r'G:\supertonic\assets\onnx'
opts = ort.SessionOptions()
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
t0 = time.time()
try:
    cfgs = load_cfgs(onnx_dir)
    dp, te, ve, vo = load_onnx_all(onnx_dir, opts, providers)
    print("Session providers dp:", dp.get_providers())
    print("Session providers vocoder:", vo.get_providers())
    tp = load_text_processor(onnx_dir)
    tts = TextToSpeech(cfgs, tp, dp, te, ve, vo)
    print("LOAD_SECONDS=%.2f" % (time.time() - t0))
    style = load_voice_style([r'G:\supertonic\assets\voice_styles\F2.json'], verbose=True)
    text = "Hai hai! Aku senang banget hari ini, ayo main bareng yuk! Hihihi, seru banget kan?"
    with timer("Generating speech from text (GPU)"):
        wav, duration = tts(text, "id", style, 20, 1.0)
    print("DURATIONS:", duration)
    outdir = r'G:\supertonic\py\results_k1gpu'
    os.makedirs(outdir, exist_ok=True)
    out = os.path.join(outdir, "Hai_hai__Aku_senang__1.wav")
    w = wav[0, : int(tts.sample_rate * duration[0].item())]
    sf.write(out, w, tts.sample_rate)
    print("Saved:", out, "sample_rate:", tts.sample_rate, "samples:", len(w))
    print("TOTAL_WALL_SECONDS=%.2f" % (time.time() - t0))
except Exception as e:
    import traceback
    print("GPU TEST FAILED:", type(e).__name__, str(e))
    traceback.print_exc()
    print("TOTAL_WALL_SECONDS=%.2f" % (time.time() - t0))