import os
import subprocess
import sys

def generate_encoder_decoder_onnx(checkpoint_path, output_dir):
    result = subprocess.run(
        [sys.executable, "-m", "piper_train.export_onnx_streaming", checkpoint_path, output_dir],
        capture_output=True,
        text=True,
        check=False
    )
    return result


def search_for_models_ckpt(base_dir):
    results = []

    for d1 in os.listdir(base_dir):
        if d1.startswith('.'):
            continue
        p1 = os.path.join(base_dir, d1)
        if not os.path.isdir(p1):
            continue

        for d2 in os.listdir(p1):
            if d2.startswith('.'):
                continue
            p2 = os.path.join(p1, d2)
            if not os.path.isdir(p2):
                continue

            for d3 in os.listdir(p2):
                if d3.startswith('.'):
                    continue
                p3 = os.path.join(p2, d3)
                if not os.path.isdir(p3):
                    continue

                for d4 in os.listdir(p3):
                    if d4.startswith('.'):
                        continue
                    p4 = os.path.join(p3, d4)
                    if not os.path.isdir(p4):
                        continue

                    for filename in os.listdir(p4):
                        if filename.startswith('.'):
                            continue
                        if filename.endswith(".ckpt"):
                            full_path = os.path.join(p4, filename)
                            results.append((d1, d2, d3, d4, filename, full_path,p4 ))

    return results

# Search for checkpoints of models to convert
ckpt_entries = search_for_models_ckpt(".")

for model_checkpoint in ckpt_entries:
    print(f"\nExporting to ONNX the encoder and decoder of PiperTTS voice from checkpoint: {model_checkpoint[5]}")
    generate_encoder_decoder_onnx(model_checkpoint[5],model_checkpoint[6])
    print(f"Generated encoder.onnx and decoder.onnx in directory: {model_checkpoint[6]}\n\n")