using System;
using System.IO;
using System.Text;
using System.Threading.Tasks;
using Microsoft.ML.OnnxRuntimeGenAI;
using UnityEngine;

namespace Sky.OnnxRuntime.Samples
{
    /// <summary>
    /// Generates text with ONNX Runtime GenAI (onnxruntime-genai).
    ///
    /// Unlike the other samples, this one needs a real GenAI model, which is far
    /// too large to embed. Point <see cref="_modelFolderPath"/> at a folder that
    /// contains an ONNX Runtime GenAI model (the folder with
    /// <c>genai_config.json</c>, the tokenizer files and the <c>*.onnx</c>
    /// weights). You can download ready-to-use models from Hugging Face, e.g.
    /// <c>microsoft/Phi-3.5-mini-instruct-onnx</c> (pick a CPU/int4 variant).
    /// </summary>
    /// <remarks>
    /// Generation runs on a background thread so it does not freeze the Editor.
    /// </remarks>
    public sealed class GenAITextGenerationSample : MonoBehaviour
    {
        [Tooltip("Absolute path to a folder containing an onnxruntime-genai model " +
                 "(genai_config.json, tokenizer files and the .onnx weights).")]
        [SerializeField]
        private string _modelFolderPath = "";

        [TextArea]
        [SerializeField]
        private string _prompt = "Write a haiku about the Unity game engine.";

        [Tooltip("Maximum number of tokens in the prompt + response.")]
        [SerializeField]
        private int _maxLength = 256;

        private void Start()
        {
            if (string.IsNullOrWhiteSpace(_modelFolderPath))
            {
                Debug.LogWarning(
                    "[ORT GenAI] Set 'Model Folder Path' on the GenAITextGenerationSample " +
                    "component to a local onnxruntime-genai model folder, then enter Play Mode.");
                return;
            }

            if (!Directory.Exists(_modelFolderPath))
            {
                Debug.LogError($"[ORT GenAI] Model folder not found: {_modelFolderPath}");
                return;
            }

            string modelPath = _modelFolderPath;
            string prompt = _prompt;
            int maxLength = _maxLength;

            // Run off the main thread — model loading and generation are blocking.
            Task.Run(() =>
            {
                try
                {
                    Generate(modelPath, prompt, maxLength);
                }
                catch (Exception e)
                {
                    Debug.LogError($"[ORT GenAI] Generation failed: {e}");
                }
            });
        }

        private static void Generate(string modelPath, string prompt, int maxLength)
        {
            Debug.Log($"[ORT GenAI] Loading model from: {modelPath}");

            using var model = new Model(modelPath);
            using var tokenizer = new Tokenizer(model);

            // Encode the prompt. Chat/instruct models usually expect a template;
            // apply it here if your model ships one.
            using var sequences = tokenizer.Encode(prompt);

            using var generatorParams = new GeneratorParams(model);
            generatorParams.SetSearchOption("max_length", maxLength);

            using var generator = new Generator(model, generatorParams);
            generator.AppendTokenSequences(sequences);

            // Stream tokens as they are produced.
            using var tokenizerStream = tokenizer.CreateStream();
            var response = new StringBuilder();

            while (!generator.IsDone())
            {
                generator.GenerateNextToken();

                ReadOnlySpan<int> newTokens = generator.GetNextTokens();
                if (newTokens.Length == 0)
                    continue;

                string piece = tokenizerStream.Decode(newTokens[newTokens.Length - 1]);
                response.Append(piece);
            }

            Debug.Log($"[ORT GenAI] Prompt:\n{prompt}");
            Debug.Log($"[ORT GenAI] Response:\n{response}");
            Debug.Log("[ORT GenAI] Generation complete ✅");
        }
    }
}