﻿#if UNITY_EDITOR
using System;
using System.Collections.Generic;
using System.IO;
using UnityEditor;
using UnityEditor.SceneManagement;
using UnityEngine;
using UnityEngine.Assertions;
using UnityEngine.Rendering;
using UnityEngine.SceneManagement;
using Object = UnityEngine.Object;

namespace ModestTree
{
    [InitializeOnLoad]
    public static class ExportHelper
    {
        static readonly List<Component> _componentsList;
        static readonly Dictionary<Material, Material> _modifiedMaterials;
        static readonly Dictionary<Type, Dictionary<AbScenePlatforms, Action<object>>> _componentCleanupFunctions;

        static readonly AbScenePlatforms[] _uniquePlatforms = new AbScenePlatforms[]
        {
            AbScenePlatforms.Pc,
            AbScenePlatforms.Android,
            AbScenePlatforms.Ios,
            AbScenePlatforms.WebGl,
            AbScenePlatforms.OculusQuest,
        };

        static ExportHelper()
        {
            Assert.IsNull(_componentCleanupFunctions);

            _componentsList = new List<Component>();
            _modifiedMaterials = new Dictionary<Material, Material>();
            _componentCleanupFunctions = new Dictionary<Type, Dictionary<AbScenePlatforms, Action<object>>>();

            RegisterComponentCleanupHandler<Camera>(AbScenePlatforms.All, RunComponentCleanup);
            RegisterComponentCleanupHandler<ParticleSystemRenderer>(AbScenePlatforms.OculusQuest, RunParticleSystemCleanupLowQuality);
        }

        static void BuildBundleFromInfo(ExportInfo info)
        {
            EditorUtility.DisplayProgressBar("Exporting Bundle(s)", "Please wait, Exporting bundles...", 0f);
            CleanupOutputAssetDirectory();

            for (int i = 0; i < info.Platforms.Length; i++)
            {
                var platform = info.Platforms[i];
                var progress = i / info.Platforms.Length;

                if (platform == AbScenePlatforms.Android || platform == AbScenePlatforms.OculusQuest)
                {
                    EditorUtility.DisplayProgressBar("Exporting Bundle(s)",
                        $"Please wait, Setting graphics API for {platform}...", progress);

                    switch (platform)
                    {
                        case AbScenePlatforms.Android:
                            ExportProjectConfig.SetGraphicsApiToOpenGLES3();
                            break;
                        case AbScenePlatforms.OculusQuest:
                            ExportProjectConfig.SetGraphicsApiToVulkan();
                            break;
                        default:
                            throw new NotSupportedException($"Unexpected platform {platform}");
                    }
                }

                EditorUtility.DisplayProgressBar("Exporting Bundle(s)",
                    $"Please wait, Exporting bundle for {platform}...", progress);

                BuildAssetBundle(info.Options, platform, info.ExportPath, info.AssetPaths);
            }

            EditorUtility.ClearProgressBar();
        }

        [MenuItem("Assets/Modest Tree Export/All Platforms/Single Model Compressed")]
        public static void BuildAllAssetBundlesCompressedAllPlatforms()
        {
            var outputPath = ShowSelectSaveDirectory();

            if (string.IsNullOrEmpty(outputPath))
            {
                return;
            }

            var assetPath = GetFullObjectPath(Selection.activeObject);

            BuildBundleFromInfo(new ExportInfo
            {
                Platforms = _uniquePlatforms,
                Options = BuildAssetBundleOptions.None,
                AssetPaths = new string[] { assetPath },
                ExportPath = outputPath
            });
        }

        [MenuItem("Assets/Modest Tree Export/All Platforms/Single Model Uncompressed")]
        public static void BuildAllAssetBundlesUncompressedAllPlatforms()
        {
            var outputPath = ShowSelectSaveDirectory();

            if (string.IsNullOrEmpty(outputPath))
            {
                return;
            }

            var assetPath = GetFullObjectPath(Selection.activeObject);

            BuildBundleFromInfo(new ExportInfo
            {
                Platforms = _uniquePlatforms,
                Options = BuildAssetBundleOptions.UncompressedAssetBundle,
                AssetPaths = new string[] { assetPath },
                ExportPath = outputPath
            });
        }

        public static void ExportBundles(AbScenePlatforms[] abScenePlatforms, BuildAssetBundleOptions options)
        {
            var outputPath = ShowSelectSaveDirectory();

            if (string.IsNullOrEmpty(outputPath))
            {
                return;
            }

            var assetPaths = new List<string>();

            foreach (var selected in Selection.objects)
            {
                if (CheckSelected(selected))
                {
                    assetPaths.Add(GetFullObjectPath(selected));
                }
            }

            if (assetPaths.Count == 0)
            {
                return;
            }

            BuildBundleFromInfo(new ExportInfo
            {
                Platforms = abScenePlatforms,
                Options = options,
                AssetPaths = assetPaths.ToArray(),
                ExportPath = outputPath
            });
        }

        [MenuItem("Assets/Modest Tree Export/Windows/Single Model Compressed")]
        public static void BuildAllAssetBundlesCompressedWindows()
        {
            BuildAssetBundleSingle(BuildAssetBundleOptions.None, AbScenePlatforms.Pc);
        }

        [MenuItem("Assets/Modest Tree Export/Windows/Single Model Uncompressed")]
        public static void BuildAllAssetBundlesUncompressedWindows()
        {
            BuildAssetBundleSingle(BuildAssetBundleOptions.UncompressedAssetBundle, AbScenePlatforms.Pc);
        }

        [MenuItem("Assets/Modest Tree Export/Android/Single Model Compressed")]
        public static void BuildAllAssetBundlesCompressedAndroid()
        {
            BuildAssetBundleSingle(BuildAssetBundleOptions.None, AbScenePlatforms.Android);
        }

        [MenuItem("Assets/Modest Tree Export/Android/Single Model Uncompressed")]
        public static void BuildAllAssetBundlesUncompressedAndroid()
        {
            BuildAssetBundleSingle(BuildAssetBundleOptions.UncompressedAssetBundle, AbScenePlatforms.Android);
        }

        [MenuItem("Assets/Modest Tree Export/OculusQuest/Single Model Compressed")]
        public static void BuildAllAssetBundlesCompressedOculusQuest()
        {
            BuildAssetBundleSingle(BuildAssetBundleOptions.None, AbScenePlatforms.OculusQuest);
        }

        [MenuItem("Assets/Modest Tree Export/OculusQuest/Single Model Uncompressed")]
        public static void BuildAllAssetBundlesUncompressedOculusQuest()
        {
            BuildAssetBundleSingle(BuildAssetBundleOptions.UncompressedAssetBundle, AbScenePlatforms.OculusQuest);
        }

        [MenuItem("Assets/Modest Tree Export/WebGl/Single Model Compressed")]
        public static void BuildAllAssetBundlesCompressedWebGl()
        {
            BuildAssetBundleSingle(BuildAssetBundleOptions.None, AbScenePlatforms.WebGl);
        }

        [MenuItem("Assets/Modest Tree Export/WebGl/Single Model Uncompressed")]
        public static void BuildAllAssetBundlesUncompressedWebGl()
        {
            BuildAssetBundleSingle(BuildAssetBundleOptions.UncompressedAssetBundle, AbScenePlatforms.WebGl);
        }

        [MenuItem("Assets/Modest Tree Export/iOS/Single Model Compressed")]
        public static void BuildAllAssetBundlesCompressediOS()
        {
            BuildAssetBundleSingle(BuildAssetBundleOptions.None, AbScenePlatforms.Ios);
        }

        [MenuItem("Assets/Modest Tree Export /iOS/Single Model Uncompressed")]
        public static void BuildAllAssetBundlesUncompressediOS()
        {
            BuildAssetBundleSingle(BuildAssetBundleOptions.UncompressedAssetBundle, AbScenePlatforms.Ios);
        }

        [MenuItem("Assets/Modest Tree Export/Current Platform/Compressed", true)]
        [MenuItem("Assets/Modest Tree Export/Current Platform/Uncompressed", true)]
        [MenuItem("Assets/Modest Tree Export/Windows/Compressed", true)]
        [MenuItem("Assets/Modest Tree Export/Windows/Uncompressed", true)]
        [MenuItem("Assets/Modest Tree Export/Android/Compressed", true)]
        [MenuItem("Assets/Modest Tree Export/Android/Uncompressed", true)]
        [MenuItem("Assets/Modest Tree Export/WebGl/Compressed", true)]
        [MenuItem("Assets/Modest Tree Export/WebGl/Uncompressed", true)]
        [MenuItem("Assets/Modest Tree Export/iOS/Compressed", true)]
        [MenuItem("Assets/Modest Tree Export/iOS/Uncompressed", true)]
        public static bool CheckSelected()
        {
            var selected = Selection.activeObject;
            return CheckSelected(selected);
        }

        public static bool CheckSelected(Object selected)
        {
            if (selected == null)
            {
                return false;
            }

            var path = GetFullObjectPath(selected);
            return GetResultExportType(path) != ExportTypes.Invalid;
        }

        static string ShowSelectSaveDirectory()
        {
            var lastDirectory = EditorPrefs.GetString("LastDir");
            var outputPath = EditorUtility.SaveFolderPanel("Select directory to export", lastDirectory, null);
            if (string.IsNullOrEmpty(outputPath))
            {
                return outputPath;
            }

            EditorPrefs.SetString("LastDir", outputPath);
            return outputPath;
        }

        static string GetFullPathForAssetPath(string assetPath)
        {
            return Path.GetFullPath(Path.Combine(Application.dataPath, "..", assetPath));
        }

        static string GetObjectPath(Object @object)
        {
            return AssetDatabase.GetAssetPath(@object.GetInstanceID());
        }

        static string GetFullObjectPath(Object @object)
        {
            return GetFullPathForAssetPath(AssetDatabase.GetAssetPath(@object.GetInstanceID()));
        }

        static ExportTypes GetResultExportType(string inputPath)
        {
            var extension = Path.GetExtension(inputPath);

            if (extension.Equals(".unity", StringComparison.OrdinalIgnoreCase))
            {
                return ExportTypes.AbScene;
            }

            if (extension.Equals(".fbx", StringComparison.OrdinalIgnoreCase) ||
                extension.Equals(".abfbx", StringComparison.OrdinalIgnoreCase) ||
                extension.Equals(".prefab", StringComparison.OrdinalIgnoreCase))
            {
                return ExportTypes.AbModel;
            }

            return ExportTypes.Invalid;
        }

        static void PrepareModel(string inputPath, AbScenePlatforms platform)
        {
            if (GetResultExportType(inputPath) == ExportTypes.AbScene)
            {
                var scene = EditorSceneManager.OpenScene(inputPath);
                foreach (var go in scene.GetRootGameObjects())
                {
                    PrepareObjectsRecursive(go, platform);
                }
            }
            else
            {
                var relativePath = ExportUtils.GetRelativeAssetPath(inputPath);
                var go = AssetDatabase.LoadAssetAtPath<GameObject>(relativePath);

                PrepareObjectsRecursive(go, platform);
            }
        }

        static void RunComponentCleanup(Camera camera)
        {
            camera.gameObject.tag = "EditorOnly";
        }

        static void RestoreModifiedMaterials()
        {
            foreach (var cur in _modifiedMaterials)
            {
                // Restore the source from the backup
                cur.Key.CopyPropertiesFromMaterial(cur.Value);
                EditorUtility.SetDirty(cur.Key);

                // Destroy the backup
                GameObject.DestroyImmediate(cur.Value);
            }

            _modifiedMaterials.Clear();
            AssetDatabase.SaveAssets();
        }

        static void SaveModifiedMaterials()
        {
            foreach (var cur in _modifiedMaterials)
            {
                EditorUtility.SetDirty(cur.Key);
            }

            AssetDatabase.SaveAssets();
        }

        static Material GetRestorableMaterial(Material source)
        {
            if (!_modifiedMaterials.ContainsKey(source))
            {
                // Copy settings so they can be restored later
                var backupMaterial = new Material(source);

                // Note (Aaron H): Unity 6 performs domain reloads on scene change. This causes
                // temporary materials to be garbage collected. Setting the hide flags here
                // appears to prevent the reference from being lost so that we can restore it later.
                backupMaterial.hideFlags = HideFlags.DontSave;

                // One per source material
                _modifiedMaterials.Add(source, backupMaterial);
            }

            return source;
        }

        static void RunParticleSystemCleanupLowQuality(ParticleSystemRenderer particleSystem)
        {
            foreach (var cur in particleSystem.sharedMaterials)
            {
                var mat = GetRestorableMaterial(cur);

                // Disable distortion
                mat.SetFloat("_DistortionEnabled", 0f);
                mat.DisableKeyword("EFFECT_BUMP");
                mat.SetShaderPassEnabled("ALWAYS", false);
            }
        }

        static void RegisterComponentCleanupHandler<T>(AbScenePlatforms flags, Action<T> handler)
            where T : class
        {
            if (flags == AbScenePlatforms.None)
            {
                return;
            }

            if (!_componentCleanupFunctions.TryGetValue(typeof(T), out var dict))
            {
                dict = new Dictionary<AbScenePlatforms, Action<object>>();
                _componentCleanupFunctions.Add(typeof(T), dict);
            }

            foreach (var cur in _uniquePlatforms)
            {
                if ((flags & cur) != AbScenePlatforms.None)
                {
                    if (dict.ContainsKey(cur))
                    {
                        dict[cur] += param => handler((T)param);
                    }
                    else
                    {
                        dict[cur] = param => handler((T)param);
                    }
                }
            }
        }

        static void PrepareObjectsRecursive(GameObject go, AbScenePlatforms platform)
        {
            try
            {
                foreach (var curComponentType in _componentCleanupFunctions)
                {
                    _componentsList.Clear();

                    go.GetComponents(curComponentType.Key, _componentsList);

                    if (curComponentType.Value.TryGetValue(platform, out var actionList))
                    {
                        foreach (var curComponent in _componentsList)
                        {
                            actionList(curComponent);
                        }
                    }
                }
            }
            finally
            {
                _componentsList.Clear();
            }

            // Recursively call children
            for (var i = 0; i < go.transform.childCount; i++)
            {
                PrepareObjectsRecursive(go.transform.GetChild(i).gameObject, platform);
            }
        }

        static void CleanupOutputAssetDirectory()
        {
            var outputDirectoryPath = GetFullPathForAssetPath(ExportUtils.OutputDirectoryName);
            if (Directory.Exists(outputDirectoryPath))
            {
                Directory.Delete(outputDirectoryPath, true);
            }

            Directory.CreateDirectory(outputDirectoryPath);
        }

        static void BuildAssetBundleSingle(BuildAssetBundleOptions options, AbScenePlatforms buildTarget)
        {
            var outputPath = ShowSelectSaveDirectory();

            if (string.IsNullOrEmpty(outputPath))
            {
                return;
            }

            var assetPath = GetFullObjectPath(Selection.activeObject);

            BuildBundleFromInfo(new ExportInfo
            {
                Platforms = new AbScenePlatforms[] { buildTarget },
                Options = options,
                AssetPaths = new string[] { assetPath },
                ExportPath = outputPath
            });
        }

        static void BuildAssetBundle(
            BuildAssetBundleOptions options,
            AbScenePlatforms buildTarget,
            string outputPath,
            string[] assetPaths)
        {
            var outputDirectoryPath = GetFullPathForAssetPath(ExportUtils.OutputDirectoryName);

            foreach (var assetPath in assetPaths)
            {
                var exportType = GetResultExportType(assetPath);

                if (exportType == ExportTypes.Invalid)
                {
                    Debug.LogError($"Invalid input file: {assetPath}");
                    return;
                }

                Debug.Log($"Exporting selected object ({exportType}): '{assetPath}'");

                if (!BuildAssetBundleInternal(outputDirectoryPath, assetPath, options, buildTarget))
                {
                    return;
                }

                CopyBuildToDestination(buildTarget, outputPath, outputDirectoryPath, assetPath, exportType);
            }

            Debug.Log("Exporting completed!");
        }

        static string GenerateFileName(string path, AbScenePlatforms target, ExportTypes exportType)
        {
            return $"{Path.GetFileNameWithoutExtension(path)}_{target}.{exportType}";
        }

        static void CopyBuildToDestination(
            AbScenePlatforms buildTarget,
            string outputPath,
            string outputDirectoryPath,
            string path,
            ExportTypes exportType)
        {
            foreach (var file in Directory.GetFiles(outputDirectoryPath))
            {
                var targetFileName = GenerateFileName(path, buildTarget, exportType);
                if (Path.GetFileName(file).Equals(targetFileName, StringComparison.OrdinalIgnoreCase))
                {
                    ProcessAssetBundle(Path.GetFullPath(file), outputPath);
                }
            }
        }

        static bool BuildAssetBundleInternal(
            string outputDirectoryPath,
            string inputPath,
            BuildAssetBundleOptions options,
            AbScenePlatforms buildPlatform)
        {
            var resultExportType = GetResultExportType(inputPath);
            var assetName = GenerateFileName(inputPath, buildPlatform, resultExportType);
            var buildName = assetName.Replace(ExportUtils.TempModelName, string.Empty);

            var outputPath = Path.Combine(outputDirectoryPath, buildName);
            var sceneManagerSetup = EditorSceneManager.GetSceneManagerSetup();

            try
            {
                PrepareModel(inputPath, buildPlatform);
                SaveModifiedMaterials();

                var subPath = inputPath.Substring(Application.dataPath.Length + 1).Replace('\\', '/');
                var assetPath = "Assets/" + subPath;

                Debug.Log($"AssetBundle building \"({buildPlatform}):{buildName}\" from \"{assetPath}\" to \"{outputPath}\"");

                var assetBundleInfo = new AssetBundleBuild
                {
                    assetBundleName = buildName
                };

                string path = null;
                string renderAssetTypeName = null;

                // Currently this has an issue where since its using the same name for configuration scene,
                // asset manager will fail to load multiple exported abscenes. 
                // HACK: Disabled until we have URP support and this is required
                if (false && resultExportType == ExportTypes.AbScene)
                {
                    WriteConfigurationAssetScene(out path, out renderAssetTypeName);
                }

                AssetDatabase.Refresh();

                var assetsList = new List<string>(new string[] { assetPath });
                if (!string.IsNullOrEmpty(path))
                {
                    assetsList.Add(path);
                }

                assetBundleInfo.assetNames = assetsList.ToArray();

                var buildTarget = GetBuildTargetFromPlatform(buildPlatform);
                var buildList = new AssetBundleBuild[] { assetBundleInfo };

                var additionalOptions = BuildAssetBundleOptions.ForceRebuildAssetBundle;
                options |= additionalOptions;

                var assetBundleManifest = BuildPipeline
                    .BuildAssetBundles(ExportUtils.OutputDirectoryName, buildList, options, buildTarget);

                if (assetBundleManifest == null)
                {
                    Debug.LogError("Failed to build asset bundle archive");
                    return false;
                }
                else
                {
                    var compressed = !options.HasFlag(BuildAssetBundleOptions.UncompressedAssetBundle);
                    if (!AbSceneExtensionHelper.WriteAssetBundleExportMetadataToFile(
                            outputPath, renderAssetTypeName, buildPlatform, compressed))
                    {
                        Debug.LogError("Failed to write additional metadata to the asset bundle!");
                        return false;
                    }
                }
            }
            catch (Exception e)
            {
                Debug.LogError($"Error building asset bundle archive:\n{e}");
                return false;
            }
            finally
            {
                RestoreModifiedMaterials();

                if (sceneManagerSetup.Length > 0)
                {
                    EditorSceneManager.RestoreSceneManagerSetup(sceneManagerSetup);
                }
            }

            Debug.Log($"Export finished, {GetResultExportType(inputPath)} saved to \"{outputPath}\"");

            return true;
        }

        static void WriteConfigurationAssetScene(out string path, out string renderAssetTypeName)
        {
            var newScene = EditorSceneManager.NewScene(NewSceneSetup.EmptyScene, NewSceneMode.Single);
            newScene.name = ExportConfigData.ConfigurationSceneName;
            path = $"Assets/{newScene.name}.unity";
            var qualitySettingsPipeline = QualitySettings.renderPipeline;
            var defaultSettingsPipeline = GraphicsSettings.currentRenderPipeline;

            RenderPipelineAsset renderAsset = null;
            renderAssetTypeName = ExportConfigData.InBuiltPipelineName;
            if (qualitySettingsPipeline != null)
            {
                renderAssetTypeName = qualitySettingsPipeline.GetType().Name;
                renderAsset = qualitySettingsPipeline;
            }
            else if (defaultSettingsPipeline != null)
            {
                renderAssetTypeName = defaultSettingsPipeline.GetType().Name;
                renderAsset = defaultSettingsPipeline;
            }

            var gameObject = new GameObject(ExportConfigData.ConfigurationGameObjectName);
            var dataObject = gameObject.AddComponent<ExportConfigData>();
            dataObject.RenderAsset = renderAsset;

            SceneManager.MoveGameObjectToScene(gameObject, newScene);
            EditorSceneManager.SaveScene(newScene, path);
        }

        private static BuildTarget GetBuildTargetFromPlatform(AbScenePlatforms platform)
        {
            BuildTarget buildTarget;
            switch (platform)
            {
                case AbScenePlatforms.Pc:
                    buildTarget = BuildTarget.StandaloneWindows64;
                    break;

                case AbScenePlatforms.Android:
                case AbScenePlatforms.OculusQuest:
                    buildTarget = BuildTarget.Android;
                    break;

                case AbScenePlatforms.WebGl:
                    buildTarget = BuildTarget.WebGL;
                    break;

                case AbScenePlatforms.Ios:
                    buildTarget = BuildTarget.iOS;
                    break;

                default:
                    throw new NotSupportedException("Unknown pipeline platform");
            }

            return buildTarget;
        }

        static void ProcessAssetBundle(string assetBundlePath, string outputDir)
        {
            var manifestFilePath = assetBundlePath + ".manifest";

            if (!File.Exists(manifestFilePath))
            {
                throw new Exception(
                    $"Expected manifest file at '{manifestFilePath}' not found");
            }

            var outputPath = Path.GetFullPath(Path.Combine(outputDir, Path.GetFileName(assetBundlePath)));

            if (outputPath.Equals(assetBundlePath))
            {
                return;
            }

            Debug.Log($"Copying '{assetBundlePath}' to '{outputPath}'");
            File.Copy(assetBundlePath, outputPath, true);
        }
    }

    public class ExportInfo
    {
        public AbScenePlatforms[] Platforms;
        public BuildAssetBundleOptions Options;
        public string[] AssetPaths;
        public string ExportPath;
    }

    enum ExportTypes
    {
        Invalid,
        AbScene,
        AbModel
    }
}
#endif
