import argparse
import os
import shutil

from tqdm import tqdm

from ort_common import WenetONNXRunner, pack_calibration_dataset


def get_args():
    parser = argparse.ArgumentParser(
        description="Generate calibration_dataset for exported ONNX models")
    parser.add_argument("--input",
                        "-i",
                        nargs="+",
                        required=True,
                        help="Input wav file(s) or directory/directories")
    parser.add_argument("--config",
                        required=True,
                        help="yaml file in checkpoint path")
    parser.add_argument(
        "--vocab",
        required=True,
        help="pretrained units.txt, for example pretrained/<model>/units.txt",
    )
    parser.add_argument("--onnx_dir",
                        default="onnx_model",
                        help="directory containing exported ONNX models")
    parser.add_argument("--calib_data_path",
                        default="calibration_dataset",
                        help="output calibration dataset directory")
    parser.add_argument("--parts",
                        nargs="+",
                        choices=["all", "offline", "online", "decoder"],
                        default=["all"],
                        help="which model inputs to generate")
    parser.add_argument("--offline_seq_len", type=int, default=1024)
    parser.add_argument("--decoder_len", type=int, default=32)
    parser.add_argument("--decoding_chunk_size", type=int, default=16)
    parser.add_argument("--num_decoding_left_chunks", type=int, default=5)
    parser.add_argument("--max_num",
                        type=int,
                        default=100,
                        help="maximum number of audio files used for calibration; set <= 0 to use all")
    parser.add_argument("--keep_existing",
                        action="store_true",
                        help="append to an existing calibration directory")
    return parser.parse_args()


def expand_audio_inputs(inputs):
    audio_exts = {".wav", ".flac", ".mp3", ".m4a", ".ogg"}
    audio_files = []
    for path in inputs:
        if os.path.isdir(path):
            for root, _, files in os.walk(path):
                for filename in files:
                    if os.path.splitext(filename)[1].lower() in audio_exts:
                        audio_files.append(os.path.join(root, filename))
        else:
            audio_files.append(path)
    audio_files = sorted(audio_files)
    if not audio_files:
        raise FileNotFoundError("No audio files found")
    return audio_files


def normalize_parts(parts):
    if "all" in parts:
        return {"offline", "online", "decoder"}
    return set(parts)


def limit_audio_files(audio_files, max_num):
    if max_num is None or max_num <= 0:
        return audio_files
    return audio_files[:max_num]


def main():
    args = get_args()
    parts = normalize_parts(args.parts)
    audio_files = limit_audio_files(expand_audio_inputs(args.input),
                                    args.max_num)

    if os.path.exists(args.calib_data_path) and not args.keep_existing:
        shutil.rmtree(args.calib_data_path)
    os.makedirs(args.calib_data_path, exist_ok=True)

    runner = WenetONNXRunner(
        args.config,
        args.vocab,
        onnx_dir=args.onnx_dir,
        offline_seq_len=args.offline_seq_len,
        decoder_len=args.decoder_len,
        decoding_chunk_size=args.decoding_chunk_size,
        num_decoding_left_chunks=args.num_decoding_left_chunks,
    )

    counts = {"offline": 0, "online": 0, "decoder": 0}
    progress = tqdm(audio_files,
                    desc="Generating calibration data",
                    unit="wav")
    for audio_idx, audio_file in enumerate(progress):
        sample_counts = runner.save_calibration_for_audio(
            audio_file, parts, args.calib_data_path, audio_idx)
        for key, value in sample_counts.items():
            counts[key] += value
        progress.set_postfix(offline=counts["offline"],
                             online=counts["online"],
                             decoder=counts["decoder"])

    print("Packing calibration dataset...")
    pack_calibration_dataset(args.calib_data_path)
    print(f"Generated calibration data in {args.calib_data_path}")
    print(f"offline samples: {counts['offline']}")
    print(f"online samples: {counts['online']}")
    print(f"decoder samples: {counts['decoder']}")


if __name__ == "__main__":
    main()