using System;
using System.Collections.Generic;
using System.Linq;
using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;
using UnityEngine;

namespace Sky.OnnxRuntime.Samples
{
    /// <summary>
    /// Demonstrates ONNX Runtime Extensions custom operators. Extensions add
    /// operators for text/vision/audio pre- and post-processing (tokenizers,
    /// string ops, image decoders, …) that are not part of the base ONNX opset.
    ///
    /// This sample enables them with a single call —
    /// <see cref="SessionOptions.RegisterOrtExtensions"/> — and then runs a tiny
    /// embedded model that uses the <c>StringUpper</c> custom op
    /// (domain <c>ai.onnx.contrib</c>) to upper-case string tensors. No external
    /// model file is required.
    /// </summary>
    public sealed class ExtensionsSample : MonoBehaviour
    {
        // A 139-byte ONNX model with a single node:
        //   text_out = StringUpper(text_in)   in domain "ai.onnx.contrib"
        // Both tensors are 1-D string tensors.
        private const string ModelBase64 =
            "CAk6bAoxCgd0ZXh0X2luEgh0ZXh0X291dCILU3RyaW5nVXBwZXI6D2FpLm9ubnguY29udHJpYhIMc3RyaW5nX3VwcGVyWhMKB3RleHRfaW4SCAoGCAgSAgoAYhQKCHRleHRfb3V0EggKBggIEgIKAEIECgAQDUITCg9haS5vbm54LmNvbnRyaWIQAQ==";

        [SerializeField]
        private string[] _inputs = { "hello unity", "OnnxRuntime" };

        private void Start()
        {
            RunStringUpper();
        }

        /// <summary>Runs the embedded StringUpper model and logs the output.</summary>
        public void RunStringUpper()
        {
            try
            {
                byte[] modelBytes = Convert.FromBase64String(ModelBase64);

                using var options = new SessionOptions();

                // Registers all ONNX Runtime Extensions custom ops with this session.
                // The extensions native library is shipped with this package.
                options.RegisterOrtExtensions();

                using var session = new InferenceSession(modelBytes, options);

                var inputTensor = new DenseTensor<string>(_inputs, new int[] { _inputs.Length });
                var inputs = new List<NamedOnnxValue>
                {
                    NamedOnnxValue.CreateFromTensor("text_in", inputTensor)
                };

                using var results = session.Run(inputs);
                string[] output = results.First().AsTensor<string>().ToArray();

                Debug.Log($"[ORT Extensions] input  = [{string.Join(", ", _inputs)}]");
                Debug.Log($"[ORT Extensions] output = [{string.Join(", ", output)}]");

                bool ok = output.Zip(_inputs, (o, i) => o == i.ToUpperInvariant()).All(x => x);
                if (ok)
                    Debug.Log("[ORT Extensions] StringUpper custom op succeeded ✅");
                else
                    Debug.LogError("[ORT Extensions] Unexpected output ❌");
            }
            catch (Exception e)
            {
                Debug.LogError(
                    "[ORT Extensions] Failed to run the extensions custom op. Make sure the " +
                    "ONNX Runtime Extensions native library is present for this platform.\n" + e);
            }
        }
    }
}