// Copyright (c) Microsoft Corporation. 
// Licensed under the MIT license. 

using SimplygonUnityEditorPlugin.Logger;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using Unity.Formats.USD;
using UnityEditor;
using UnityEngine;
using USD.NET;

namespace Simplygon.Unity.EditorPlugin
{
    public class SimplygonImporter
    {
        public static void Import(ISimplygon simplygon, spPipeline simplygonPipeline, ref int lodIndex, string assetFolderPath, string gameObjectName, List<GameObject> importedGameObjects = null)
        {
            if (!simplygonPipeline.GetPipelineSettings().GetIntermediateStep())
            {
                using (spScene simplygonScene = simplygonPipeline.GetProcessedScene())
                {
                    string assetLODFolderGuid = AssetDatabase.CreateFolder(assetFolderPath, $"LOD{lodIndex}");
                    string assetLODFolderPath = AssetDatabase.GUIDToAssetPath(assetLODFolderGuid);

                    string lodPrefabName = $"{gameObjectName}_LOD{lodIndex}";
                    string projectPath = System.IO.Path.GetFullPath(System.IO.Path.Combine(Application.dataPath, "../"));
                    string usdFilePath = Path.Combine(projectPath, assetLODFolderPath, lodPrefabName + ".usd");

                    using (spSceneExporter sgSceneExporter = simplygon.CreateSceneExporter())
                    {
                        sgSceneExporter.SetExportFilePath(usdFilePath);
                        sgSceneExporter.SetScene(simplygonScene);
                        sgSceneExporter.RunExport();
                    }

                    if (File.Exists(usdFilePath))
                    {
                        AssetDatabase.Refresh();
                        var go = Import(usdFilePath, assetLODFolderPath, lodPrefabName);

                        if (go != null && importedGameObjects != null)
                        {
                            importedGameObjects.Add(go);
                        }
                    }
                    else
                    {
                        SimplygonLogger.Instance.Log(Category.Error, $"Failed to process LOD{lodIndex}.");
                    }
                }
                lodIndex++;
            }

            for (uint i = 0; i < simplygonPipeline.GetCascadedPipelineCount(); ++i)
            {
                Import(simplygon, simplygonPipeline.GetCascadedPipelineByIndex(i), ref lodIndex, assetFolderPath, gameObjectName, importedGameObjects);
            }
        }

        public static GameObject Import(string usdFilePath, string assetLODFolderPath, string gameObjectName)
        {
            InitUsd.Initialize();
            var stage = pxr.UsdStage.Open(usdFilePath, pxr.UsdStage.InitialLoadSet.LoadAll);
            var scene = Scene.Open(stage);

            scene.Time = 1.0;

            var importOptions = new SceneImportOptions();
            importOptions.projectAssetPath = assetLODFolderPath + "/";
            importOptions.changeHandedness = BasisTransformation.SlowAndSafe;
            importOptions.materialImportMode = MaterialImportMode.ImportPreviewSurface;
            importOptions.payloadPolicy = PayloadPolicy.LoadAll;
            importOptions.usdRootPath = GetDefaultRoot(scene);
            importOptions.meshOptions.normals = ImportMode.Import;
            importOptions.meshOptions.tangents = ImportMode.Import;
            importOptions.meshOptions.color = ImportMode.Import;
            importOptions.meshOptions.texcoord0 = ImportMode.Import;
            importOptions.meshOptions.texcoord1 = ImportMode.Import;
            importOptions.meshOptions.texcoord2 = ImportMode.Import;
            importOptions.meshOptions.texcoord3 = ImportMode.Import;
            importOptions.materialMap.useOriginalShaderIfAvailable = true;

            var go = new GameObject(gameObjectName);
            SceneImporter.ImportUsd(go, scene, new PrimMap(), importOptions);

            scene.Close();

            var skinnedMeshRenderers = go.GetComponentsInChildren<SkinnedMeshRenderer>();
            var meshRenderers = go.GetComponentsInChildren<MeshRenderer>();
            var meshFilters = go.GetComponentsInChildren<MeshFilter>();

            List<Material> materials = new List<Material>();
            List<Mesh> meshes = new List<Mesh>();

            foreach (var skinnedMeshRenderer in skinnedMeshRenderers)
            {
                if (skinnedMeshRenderer.sharedMaterials.Length > 0)
                {
                    materials.AddRange(skinnedMeshRenderer.sharedMaterials);
                }

                if (skinnedMeshRenderer.sharedMesh != null)
                {
                    meshes.Add(skinnedMeshRenderer.sharedMesh);
                }
            }

            foreach (var meshRenderer in meshRenderers)
            {
                if (meshRenderer.sharedMaterials.Length > 0)
                {
                    materials.AddRange(meshRenderer.sharedMaterials);
                }
            }

            foreach (var meshFilter in meshFilters)
            {
                if (meshFilter.sharedMesh != null)
                {
                    meshes.Add(meshFilter.sharedMesh);
                }
            }

            materials = materials.Distinct().ToList();
            meshes = meshes.Distinct().ToList();

            foreach (var material in materials)
            {
                AssetDatabase.CreateAsset(material, assetLODFolderPath + "/" + material.name.Replace("/", "_") + ".asset");

                Texture normalMap = null;
                if (material.HasProperty("_NormalMap"))
                {
                    normalMap = material.GetTexture("");
                }
                else if (material.HasProperty("_BumpMap"))
                {
                    normalMap = material.GetTexture("_BumpMap");
                }

                if (normalMap != null)
                {
                    string path = AssetDatabase.GetAssetPath(normalMap);
                    TextureImporter textureImporter = AssetImporter.GetAtPath(path) as TextureImporter;
                    textureImporter.textureType = TextureImporterType.NormalMap;
                    AssetDatabase.ImportAsset(path);
                }

                if (material.shaderKeywords.Contains("_ALPHATEST_ON")) //Cutout
                {
                    material.SetOverrideTag("RenderType", "TransparentCutout");
                    material.SetInt("_SrcBlend", (int)UnityEngine.Rendering.BlendMode.One);
                    material.SetInt("_DstBlend", (int)UnityEngine.Rendering.BlendMode.Zero);
                    material.SetInt("_ZWrite", 1);
                    material.EnableKeyword("_ALPHATEST_ON");
                    material.DisableKeyword("_ALPHABLEND_ON");
                    material.DisableKeyword("_ALPHAPREMULTIPLY_ON");
                    material.renderQueue = (int)UnityEngine.Rendering.RenderQueue.AlphaTest;
                }
                else if (material.shaderKeywords.Contains("_ALPHABLEND_ON")) //Blend
                {
                    material.SetOverrideTag("RenderType", "Transparent");
                    material.SetInt("_SrcBlend", (int)UnityEngine.Rendering.BlendMode.SrcAlpha);
                    material.SetInt("_DstBlend", (int)UnityEngine.Rendering.BlendMode.OneMinusSrcAlpha);
                    material.SetInt("_ZWrite", 0);
                    material.DisableKeyword("_ALPHATEST_ON");
                    material.EnableKeyword("_ALPHABLEND_ON");
                    material.DisableKeyword("_ALPHAPREMULTIPLY_ON");
                    material.renderQueue = (int)UnityEngine.Rendering.RenderQueue.Transparent;
                }
                else if (material.shaderKeywords.Contains("_ALPHAPREMULTIPLY_ON")) //Transparent
                {
                    material.SetOverrideTag("RenderType", "Transparent");
                    material.SetInt("_SrcBlend", (int)UnityEngine.Rendering.BlendMode.One);
                    material.SetInt("_DstBlend", (int)UnityEngine.Rendering.BlendMode.OneMinusSrcAlpha);
                    material.SetInt("_ZWrite", 0);
                    material.DisableKeyword("_ALPHATEST_ON");
                    material.DisableKeyword("_ALPHABLEND_ON");
                    material.EnableKeyword("_ALPHAPREMULTIPLY_ON");
                    material.renderQueue = (int)UnityEngine.Rendering.RenderQueue.Transparent;
                }
                //else //Opaque
                //{
                //    material.SetOverrideTag("RenderType", "");
                //    material.SetInt("_SrcBlend", (int)UnityEngine.Rendering.BlendMode.One);
                //    material.SetInt("_DstBlend", (int)UnityEngine.Rendering.BlendMode.Zero);
                //    material.SetInt("_ZWrite", 1);
                //    material.DisableKeyword("_ALPHATEST_ON");
                //    material.DisableKeyword("_ALPHABLEND_ON");
                //    material.DisableKeyword("_ALPHAPREMULTIPLY_ON");
                //    material.renderQueue = -1;
                //}
            }

            foreach (var mesh in meshes)
            {
                AssetDatabase.CreateAsset(mesh, assetLODFolderPath + "/" + mesh.name.Replace("/", "_") + ".asset");
            }

            PrefabUtility.SaveAsPrefabAsset(go, assetLODFolderPath + "/" + go.name.Replace("/", "_") + ".prefab");

            AssetDatabase.SaveAssets();

            return go;
        }

        private static string GetObjectName(pxr.SdfPath rootPrimName, string path)
        {
            return pxr.UsdCs.TfIsValidIdentifier(rootPrimName.GetName())
                 ? rootPrimName.GetName()
                 : path;
        }

        private static pxr.SdfPath GetDefaultRoot(Scene scene)
        {
            var children = scene.Stage.GetPseudoRoot().GetChildren().ToList();
            if (children.Count == 1)
            {
                return children[0].GetPath();
            }
            return pxr.SdfPath.AbsoluteRootPath();
        }
    }
}
