#!/usr/bin/env python3
"""
ONNX 人脸模型转 TFLite 脚本（基于 onnx2tf）

环境要求（推荐使用已有 onnx2tf 的 conda 环境）：
    conda activate base   # 已安装 onnx2tf + ai-edge-litert
    或 pip install onnx2tf ai-edge-litert onnx onnxruntime

用法:
    python convert_model.py <input.onnx> <output.tflite>

示例:
    python convert_model.py backup_models/buffalo_sc/w600k_mbf.onnx app/app/src/main/assets/mobilefacenet.tflite
"""

import os
import shutil
import subprocess
import sys
import tempfile


def main() -> int:
    if len(sys.argv) != 3:
        print(__doc__)
        return 1

    onnx_path = sys.argv[1]
    tflite_path = sys.argv[2]

    if not os.path.isfile(onnx_path):
        print(f"错误：找不到输入模型 {onnx_path}")
        return 1

    work_dir = tempfile.mkdtemp(prefix="onnx2tf_")
    try:
        print(f"转换 {onnx_path} -> {tflite_path}")
        subprocess.check_call(
            [sys.executable, "-m", "onnx2tf", "-i", onnx_path, "-o", work_dir]
        )

        base = os.path.splitext(os.path.basename(onnx_path))[0]
        converted = os.path.join(work_dir, f"{base}_float32.tflite")
        if not os.path.isfile(converted):
            candidates = [
                f for f in os.listdir(work_dir) if f.endswith("_float32.tflite")
            ]
            if not candidates:
                print("错误：onnx2tf 未生成 float32 tflite，请查看上方日志")
                return 1
            converted = os.path.join(work_dir, candidates[0])

        os.makedirs(os.path.dirname(tflite_path) or ".", exist_ok=True)
        shutil.copyfile(converted, tflite_path)
        print(f"完成：{tflite_path}（{os.path.getsize(tflite_path) / 1024 / 1024:.2f} MB）")
        return 0
    finally:
        shutil.rmtree(work_dir, ignore_errors=True)


if __name__ == "__main__":
    sys.exit(main())