using System.IO;
using Unity.InferenceEngine;
using UnityEditor;
using UnityEngine;

namespace Audio2Lipsync.Editor
{
    /// <summary>
    /// Turns an imported ONNX (e.g. a re-export from Tools~/export_onnx.py) into a .sentis. Put the
    /// result in StreamingAssets/Audio2Lipsync/ to use it instead of the packaged model.
    /// </summary>
    public static class Audio2LipsyncSerializer
    {
        /// <param name="quantization">null = keep fp32 weights (what the package ships).</param>
        public static void Serialize(ModelAsset onnx, string outPath, QuantizationType? quantization = null)
        {
            var model = ModelLoader.Load(onnx);
            if (quantization.HasValue) ModelQuantizer.QuantizeWeights(quantization.Value, ref model);
            Directory.CreateDirectory(Path.GetDirectoryName(outPath));
            ModelWriter.Save(outPath, model);   // reload from disk before running a quantized model
            Debug.Log($"[Audio2Lipsync] wrote {outPath} ({new FileInfo(outPath).Length / 1e6:F0} MB)");
        }

        [MenuItem("Assets/Audio2Lipsync/Serialize ONNX to StreamingAssets")]
        static void SerializeSelected()
        {
            Serialize((ModelAsset)Selection.activeObject, Audio2LipsyncModel.StreamingAssetsPath);
            AssetDatabase.Refresh();
        }

        [MenuItem("Assets/Audio2Lipsync/Serialize ONNX to StreamingAssets", true)]
        static bool SerializeSelectedValidate() => Selection.activeObject is ModelAsset;
    }
}