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>
    /// Minimal end-to-end ONNX Runtime check. Runs a tiny embedded model that
    /// computes <c>output = input * 2 + 1</c> element-wise, then verifies the
    /// result. No external model file or network access is required, so this is
    /// the fastest way to confirm the native runtime loads and runs on the
    /// current platform.
    /// </summary>
    /// <remarks>
    /// Add this component to any GameObject and enter Play Mode. Watch the
    /// Console for the inference result.
    /// </remarks>
    public sealed class BasicInferenceSample : MonoBehaviour
    {
        // A 185-byte ONNX model with two initializers:
        //   scaled = input * [2, 2, 2]
        //   output = scaled + [1, 1, 1]
        // input/output are float tensors of shape [1, 3].
        private const string ModelBase64 =
            "CAk6rgEKGwoFaW5wdXQKBXNjYWxlEgZzY2FsZWQiA011bAobCgZzY2FsZWQKBGJpYXMSBm91dHB1dCIDQWRkEgpzY2FsZV9iaWFzKhkIAxABQgVzY2FsZUoMAAAAQAAAAEAAAABAKhgIAxABQgRiaWFzSgwAAIA/AACAPwAAgD9aFwoFaW5wdXQSDgoMCAESCAoCCAEKAggDYhgKBm91dHB1dBIOCgwIARIICgIIAQoCCANCBAoAEA0=";

        private void Start()
        {
            RunInference();
        }

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

            // InferenceSession loads the ONNX Runtime native library on first use.
            using var session = new InferenceSession(modelBytes);

            var inputData = new float[] { 1f, 2f, 3f };
            var inputTensor = new DenseTensor<float>(inputData, new int[] { 1, 3 });

            var inputs = new List<NamedOnnxValue>
            {
                NamedOnnxValue.CreateFromTensor("input", inputTensor)
            };

            using IDisposableReadOnlyCollection<DisposableNamedOnnxValue> results = session.Run(inputs);

            float[] output = results.First().AsTensor<float>().ToArray();

            Debug.Log($"[ONNX Runtime] input  = [{string.Join(", ", inputData)}]");
            Debug.Log($"[ONNX Runtime] output = [{string.Join(", ", output)}]  (expected [3, 5, 7])");

            var expected = new float[] { 3f, 5f, 7f };
            bool ok = output.Length == expected.Length
                      && output.Zip(expected, (a, b) => Mathf.Abs(a - b) < 1e-4f).All(x => x);

            if (ok)
                Debug.Log("[ONNX Runtime] Basic inference succeeded ✅");
            else
                Debug.LogError("[ONNX Runtime] Basic inference produced an unexpected result ❌");
        }
    }
}