#if UNITY_EDITOR
using System;
using System.IO;
using System.Linq;
using System.Net.Http;
using UnityEditor;
using UnityEngine;
using OnDeviceAgent.Inference;

namespace OnDeviceAgent.AgentCore.Editor
{
    // The com.sky.sentis.* weights are not committed; each package declares where they live in its
    // package.json:
    //   "modelSource": { "huggingFaceRepo": "Sky-Kim/com.sky.sentis.e5", "revision": "main" }
    // This provisioner repopulates the package's Models~ folder from that Hugging Face repo so the
    // Editor (ModelPathResolver) and the build step (ProvisionToStreamingAssets) keep working
    // unchanged. It runs automatically on Editor load and before a player build; only the repo's
    // Models~/ subtree is fetched, so no per-file list needs maintaining.
    [InitializeOnLoad]
    public static class SentisModelProvisioner
    {
        const string ApiBase = "https://huggingface.co/api/models";
        const string ResolveBase = "https://huggingface.co";
        const string ModelsDir = "Models~";
        const string SkipKey = "SentisModelProvisioner.SkipThisSession";

        static SentisModelProvisioner()
        {
            // Defer past the domain reload so the load isn't blocked, then fill any missing weights.
            EditorApplication.delayCall += () =>
            {
                if (SessionState.GetBool(SkipKey, false))
                    return;
                EnsureAll();
            };
        }

        [Serializable] class ModelSource { public string huggingFaceRepo; public string revision; }
        [Serializable] class PkgManifest { public ModelSource modelSource; }
        [Serializable] class TreeEntry { public string type; public string path; public long size; }
        [Serializable] class TreeWrap { public TreeEntry[] items; }

        // Downloads the Models~ payload for every package whose folder is empty. Safe to call repeatedly:
        // packages that already have their weights are skipped without touching the network.
        public static void EnsureAll()
        {
            try
            {
                foreach (var (_, pkg) in ModelPathResolver.PackageMap)
                    EnsurePackage(pkg);
            }
            finally
            {
                EditorUtility.ClearProgressBar();
            }
        }

        static void EnsurePackage(string pkg)
        {
            var pkgRoot = Path.GetFullPath(Path.Combine("Packages", pkg));
            if (!Directory.Exists(pkgRoot))
                return;

            var modelsRoot = Path.Combine(pkgRoot, ModelsDir);
            if (HasWeights(modelsRoot))
                return; // already provisioned

            var src = ReadModelSource(pkgRoot);
            if (src == null || string.IsNullOrWhiteSpace(src.huggingFaceRepo))
            {
                Debug.LogWarning($"[SentisModelProvisioner] {pkg}: Models~ is empty and package.json has no modelSource; skipping.");
                return;
            }

            var rev = string.IsNullOrWhiteSpace(src.revision) ? "main" : src.revision.Trim();
            try
            {
                DownloadModelsSubtree(pkg, src.huggingFaceRepo.Trim(), rev, pkgRoot);
            }
            catch (OperationCanceledException)
            {
                SessionState.SetBool(SkipKey, true);
                Debug.LogWarning($"[SentisModelProvisioner] {pkg}: download canceled; will not retry this session.");
            }
            catch (Exception e)
            {
                Debug.LogError($"[SentisModelProvisioner] {pkg}: failed to download weights from {src.huggingFaceRepo}@{rev}: {e.Message}");
            }
        }

        static bool HasWeights(string modelsRoot)
        {
            if (!Directory.Exists(modelsRoot))
                return false;
            return Directory.EnumerateFiles(modelsRoot, "*", SearchOption.AllDirectories)
                            .Any(f => !f.EndsWith(".meta", StringComparison.Ordinal));
        }

        static ModelSource ReadModelSource(string pkgRoot)
        {
            var path = Path.Combine(pkgRoot, "package.json");
            if (!File.Exists(path))
                return null;
            try { return JsonUtility.FromJson<PkgManifest>(File.ReadAllText(path))?.modelSource; }
            catch (Exception) { return null; }
        }

        static void DownloadModelsSubtree(string pkg, string repo, string rev, string pkgRoot)
        {
            using var http = new HttpClient { Timeout = TimeSpan.FromMinutes(30) };

            var treeUrl = $"{ApiBase}/{repo}/tree/{rev}?recursive=true";
            var treeJson = http.GetStringAsync(treeUrl).GetAwaiter().GetResult();
            var entries = JsonUtility.FromJson<TreeWrap>("{\"items\":" + treeJson + "}").items ?? Array.Empty<TreeEntry>();

            var files = entries
                .Where(e => e.type == "file" && e.path != null &&
                            e.path.Replace('\\', '/').StartsWith(ModelsDir + "/", StringComparison.Ordinal))
                .ToArray();
            if (files.Length == 0)
                return;

            long total = files.Sum(f => Math.Max(0, f.size));
            long doneTotal = 0;
            for (int i = 0; i < files.Length; i++)
            {
                var rel = files[i].path.Replace('\\', '/');
                var dst = Path.Combine(pkgRoot, rel.Replace('/', Path.DirectorySeparatorChar));
                Directory.CreateDirectory(Path.GetDirectoryName(dst));

                var url = $"{ResolveBase}/{repo}/resolve/{rev}/{string.Join("/", rel.Split('/').Select(Uri.EscapeDataString))}";
                var label = $"{pkg}  ({i + 1}/{files.Length})  {Path.GetFileName(rel)}";
                DownloadFile(http, url, dst, label, total, ref doneTotal);
            }

            Debug.Log($"[SentisModelProvisioner] {pkg}: provisioned {files.Length} file(s) from {repo}@{rev}.");
        }

        static void DownloadFile(HttpClient http, string url, string dst, string label, long total, ref long doneTotal)
        {
            var part = dst + ".part";
            using var resp = http.GetAsync(url, HttpCompletionOption.ResponseHeadersRead).GetAwaiter().GetResult();
            resp.EnsureSuccessStatusCode();

            using (var net = resp.Content.ReadAsStreamAsync().GetAwaiter().GetResult())
            using (var file = new FileStream(part, FileMode.Create, FileAccess.Write, FileShare.None, 1 << 20))
            {
                var buf = new byte[1 << 20];
                int read;
                while ((read = net.Read(buf, 0, buf.Length)) > 0)
                {
                    file.Write(buf, 0, read);
                    doneTotal += read;
                    var frac = total > 0 ? Math.Min(1f, doneTotal / (float)total) : 0f;
                    if (EditorUtility.DisplayCancelableProgressBar("Downloading Sentis model weights", label, frac))
                        throw new OperationCanceledException();
                }
            }

            if (File.Exists(dst)) File.Delete(dst);
            File.Move(part, dst);
        }
    }
}
#endif