using System.Collections;
using System.IO;
using UnityEngine;
using UnityEngine.Networking;

namespace SentisSkills.Inference
{
    /// <summary>
    /// Where the model, tokenizer and developer block are read from. In the Editor, the package's Models~ folder in the
    /// package cache. In a player, the copy the build puts in StreamingAssets/FunctionGemmaHome (see
    /// FunctionGemmaBuildStaging). Android keeps StreamingAssets inside the APK, where files can't be opened, so
    /// <see cref="CopyOutOfApk"/> first copies them to persistentDataPath.
    /// </summary>
    public static class FunctionGemmaFiles
    {
        public const string PackageName = "com.sky-kim.functiongemma-home";
        public const string Model = "functiongemma_home_uint8.sentis";
        public const string Tokenizer = "tokenizer.json";
        public const string DevBlock = "fg_dev_block_home.txt";
        public static readonly string[] All = { Model, Tokenizer, DevBlock };

        /// <summary>StreamingAssets folder a player build copies Models~ into.</summary>
        public const string StreamingFolder = "FunctionGemmaHome";

        /// <summary>True where StreamingAssets is a URL inside the app package (Android) instead of a folder.</summary>
        static bool StreamingAssetsInPackage => Application.streamingAssetsPath.Contains("://");

        /// <summary>Folder holding <see cref="All"/>.</summary>
        public static string Root
        {
            get
            {
#if UNITY_EDITOR
                var info = UnityEditor.PackageManager.PackageInfo.FindForAssetPath($"Packages/{PackageName}");
                return info != null ? Path.Combine(info.resolvedPath, "Models~") : null;
#else
                return Path.Combine(StreamingAssetsInPackage ? Application.persistentDataPath : Application.streamingAssetsPath, StreamingFolder);
#endif
            }
        }

        public static string PathOf(string file) => Root == null ? null : Path.Combine(Root, file);

        /// <summary>
        /// Android: copies the files from the APK to persistentDataPath (once; later runs reuse the copy). Other
        /// platforms read them in place.
        /// </summary>
        public static IEnumerator CopyOutOfApk()
        {
            if (Application.isEditor || !StreamingAssetsInPackage) yield break;
            DeleteLegacyCopies();
            foreach (var file in All)
            {
                var dst = PathOf(file);
                if (File.Exists(dst)) continue;
                Directory.CreateDirectory(Path.GetDirectoryName(dst));
                var part = dst + ".part";
                using (var req = UnityWebRequest.Get($"{Application.streamingAssetsPath}/{StreamingFolder}/{file}"))
                {
                    req.downloadHandler = new DownloadHandlerFile(part) { removeFileOnAbort = true };
                    Debug.Log($"[FunctionGemma] copying {file} to {dst}");
                    yield return req.SendWebRequest();
                    if (req.result != UnityWebRequest.Result.Success)
                    {
                        Debug.LogError($"[FunctionGemma] copy of {file} failed: {req.error}");
                        yield break;
                    }
                }
                File.Move(part, dst);
            }
        }

        // Versions before 2.0 copied the files to persistentDataPath/Agent; drop them so they don't keep ~870 MB.
        static void DeleteLegacyCopies()
        {
            var legacy = Path.Combine(Application.persistentDataPath, "Agent");
            foreach (var file in new[] { Model, "functiongemma_home_f32.sentis", "functiongemma_home_uint8.sentis", Tokenizer, DevBlock })
            {
                var path = Path.Combine(legacy, file);
                if (File.Exists(path)) File.Delete(path);
            }
        }
    }
}