﻿//Collar Tool created by ArkorenArt for Spectrum avatars. Do not distribute or sell this script or any part of it, but feel free to use it as reference!

using UnityEngine;
using UnityEditor;
using UnityEditor.SceneManagement;
using UnityEditor.Presets;
using System.IO;
using System.Collections.Generic;
using System.Linq;
public class CollarTool : EditorWindow
{
    GameObject avatar = null;
    GameObject avatarQ = null;

    GameObject lastAvatar;

    Vector2 scrollPosition;

    Material collarMat;

    int collarTagDesign;
    List<string> collarTagDesigns = new List<string>();

    [SerializeField]
    Color32 collarCol;
    Color32 _collar;
    [SerializeField]
    Color32 metalCol;
    Color32 _metal;
    [SerializeField]
    Color32 bellCol;
    Color32 _bell;
    [SerializeField]
    Color32 collarFadeCol;
    Color32 _collarFade;
    bool collarFade;

    [MenuItem("Arkoren/Spectrum/Collar")]
    static void Init()
    {
        // Get existing open window or if none, make a new one:
        CollarTool window = (CollarTool)EditorWindow.GetWindow(typeof(CollarTool));
        window.Show();
    }

    void OnGUI()
    {

        EditorGUIUtility.labelWidth = 100;
        scrollPosition = GUILayout.BeginScrollView(scrollPosition);

        EditorGUILayout.Space(5);

        GUIStyle header = new GUIStyle(GUI.skin.label) { alignment = TextAnchor.MiddleCenter, fontSize = 17, fontStyle = FontStyle.Bold };
        GUIStyle header2 = new GUIStyle(GUI.skin.label) { alignment = TextAnchor.MiddleCenter, fontSize = 15 };
        GUIStyle header3 = new GUIStyle(GUI.skin.label) { alignment = TextAnchor.MiddleCenter, fontSize = 14 };
        GUIStyle header4 = new GUIStyle(GUI.skin.label) { alignment = TextAnchor.MiddleCenter, fontSize = 11 };

        EditorGUILayout.LabelField("Collar Tool", header);

        EditorGUILayout.Space(15);

        if (EditorApplication.isPlaying)
        {
            EditorGUILayout.LabelField("Editor currently in play mode.", header2);
        }
        else
        {
            if (avatar == null)
            {
                //find the avatar in the scene
                foreach (GameObject obj in EditorSceneManager.GetActiveScene().GetRootGameObjects())
                {
                    if (obj.GetComponent<Animator>() && !obj.name.Contains("Quest"))
                    {
                        if(obj.activeSelf == true)
                        {
                            lastAvatar = avatar;
                            avatar = obj;

                            break;
                        }
                    }
                }

                //find the quest avatar in the scene
                foreach (GameObject obj in EditorSceneManager.GetActiveScene().GetRootGameObjects())
                {
                    if (obj.GetComponent<Animator>() && obj.name.Contains("Quest"))
                    {
                        if (obj.activeSelf == true)
                        {
                            avatarQ = obj;

                            break;
                        }
                    }
                }
            }

            //display the avatar boxes
            avatar = (GameObject)EditorGUILayout.ObjectField("Avatar:", avatar, typeof(GameObject), true);

            avatarQ = (GameObject)EditorGUILayout.ObjectField("Quest Avatar:", avatarQ, typeof(GameObject), true);

            if (avatar == null && avatarQ == null)
            {
                EditorStyles.label.wordWrap = true;
                EditorGUILayout.LabelField("Avatars not found, make sure a scene is open and selected as the active scene, or place your avatar(s) in the boxes above.", GUILayout.ExpandHeight(false));
            }
            else if(avatar == null && avatarQ != null)
            {
                EditorStyles.label.wordWrap = true;
                EditorGUILayout.LabelField("Only Quest avatar selected, please also add a PC avatar to allow customization/quest material generation.", GUILayout.ExpandHeight(false));
            }


            if (avatar != null)
            {
                if (avatar.name.Contains("(Clone)"))
                {
                    avatar = null;
                    return;
                }

                bool collar = false;

                if (avatar.transform.Find("Armature/Hips/Spine/Chest/Neck/SpectrumCollar/Collar") != null)
                {
                    collarMat = avatar.transform.Find("Armature/Hips/Spine/Chest/Neck/SpectrumCollar/Collar").GetComponent<MeshRenderer>().sharedMaterial;

                    collar = true;
                }
                else
                {
                    EditorStyles.label.wordWrap = true;
                    EditorGUILayout.LabelField("No collar found on avatar.", GUILayout.ExpandHeight(false));

                    return;
                }

                //if first frame after inputting avatar, get current material colors and toggles
                if (lastAvatar != avatar)
                {
                    if (collarMat.HasProperty("_DecalBlendAlpha1") == false)
                    {
                        Debug.LogError("Collar material is not Poiyomi, you may have put a Quest avatar in the Avatar box.");

                        return;
                    }

                        collarFade = collarMat.GetFloat("_DecalBlendAlpha1") > 0.5f ? true : false;
                        collarCol = GetMaterialColor(collarMat, "_GlobalThemeColor0");
                        collarFadeCol = GetMaterialColor(collarMat, "_GlobalThemeColor1");
                        metalCol = GetMaterialColor(collarMat, "_GlobalThemeColor2");
                        bellCol = GetMaterialColor(collarMat, "_GlobalThemeColor3");

                    //get collar tags
                    var info = new DirectoryInfo(Application.dataPath + "/Spectrum/Textures/Collar/Tags/");
                    FileInfo[] files = info.GetFiles();

                    for (int i = 0; i < files.Length; i++)
                    {
                        if (!files[i].Name.Contains(".meta"))
                        {
                            collarTagDesigns.Add(files[i].Name.Replace(".png", ""));
                        }
                    }

                    for (int j = 0; j < collarTagDesigns.Count; j++)
                    {
                        if (avatar.transform.Find("Armature/Hips/Spine/Chest/Neck/SpectrumCollar/Tag Circle"))
                        {
                            if (avatar.transform.Find("Armature/Hips/Spine/Chest/Neck/SpectrumCollar/Tag Circle").GetComponent<MeshRenderer>().sharedMaterial.GetTexture("_DecalTexture3").name == collarTagDesigns[j])
                            {
                                collarTagDesign = j;
                            }
                        }
                    }


                    lastAvatar = avatar;
                }

                //set up variables for undo
                Color32 c14 = collarCol;
                Color32 c15 = collarFadeCol;
                Color32 c16 = metalCol;
                Color32 c17 = bellCol;

                //Begin change check
                EditorGUI.BeginChangeCheck();

                EditorGUILayout.Space(10);

                EditorGUILayout.LabelField("Customization", header2);
                EditorGUILayout.Space(2);

                //material color selection

                if (collar)
                {
                        collarFade = EditorGUILayout.Toggle(new GUIContent("Gradient fade:", "Add a color gradient to the collar."), collarFade);
                        c14 = EditorGUILayout.ColorField(new GUIContent("Collar:"), c14, true, false, true);
                        if (collarFade)
                        {
                            c15 = EditorGUILayout.ColorField(new GUIContent("Fade:"), c15, true, false, true);
                        }
                        c16 = EditorGUILayout.ColorField(new GUIContent("Metal:"), c16, true, false, true);

                        c17 = EditorGUILayout.ColorField(new GUIContent("Bell:"), c17, true, false, true);


                        collarTagDesign = EditorGUILayout.Popup(new GUIContent("Tag Design:", "Set the collar's tag design"), collarTagDesign, collarTagDesigns.ToArray());
                }

                //End change check
                if (EditorGUI.EndChangeCheck())
                {
                    //record object before changing, will undo these changes
                    Undo.RecordObject(this, "material");
                    collarCol = c14;
                    collarFadeCol = c15;
                    metalCol = c16;
                    bellCol = c17;
                }

                if (collarFade)
                {
                    collarMat.SetFloat("_DecalBlendAlpha1", 1);
                }
                else
                {
                    collarMat.SetFloat("_DecalBlendAlpha1", 0);
                }

                SetMaterialColor(collarMat, "_GlobalThemeColor0", collarCol);
                SetMaterialColor(collarMat, "_GlobalThemeColor1", collarFadeCol);
                SetMaterialColor(collarMat, "_GlobalThemeColor2", metalCol);
                SetMaterialColor(collarMat, "_GlobalThemeColor3", bellCol);
                collarMat.SetTexture("_DecalTexture3", AssetDatabase.LoadAssetAtPath<Texture2D>("Assets/Spectrum/Textures/Collar/Tags/" + collarTagDesigns[collarTagDesign] + ".png"));

            }




        if (avatar != null && avatarQ != null)
        {
            if (GUILayout.Button(new GUIContent("Generate Quest Materials", "Apply the current colors to the Quest avatar by generating new textures and materials.")))
            {
                if (AssetDatabase.IsValidFolder($"Assets/Spectrum/Generated"))
                {
                    GenerateQuestMat();
                }
                else
                {
                    string guid = AssetDatabase.CreateFolder("Assets/Spectrum", "Generated");

                    GenerateQuestMat();
                }
            }
        }
        }

        GUILayout.EndScrollView();
    }

    Color GetMaterialColor(Material mat, string property)
    {
        if (mat.HasProperty(property))
        {
            return mat.GetColor(property);
        }

        return Color.clear;
    }

    void SetMaterialColor(Material mat, string property, Color col)
    {
        if (mat.HasProperty(property))
        {
            mat.SetColor(property, col);
        }
    }

    void GenerateQuestMat()
    {
        Preset preset = AssetDatabase.LoadAssetAtPath<Preset>("Assets/Arkoren/Presets/QuestPreset.preset");


            Material collarMat = avatar.transform.Find("Armature/Hips/Spine/Chest/Neck/SpectrumCollar/Collar").GetComponent<MeshRenderer>().sharedMaterial;

        if(avatarQ.transform.Find("Armature/Hips/Spine/Chest/Neck/SpectrumCollar/Collar") == null)
        {
            Debug.LogError("Quest avatar is missing Collar", avatarQ);
            return;
        }

            Material questCollarMat = new Material(avatarQ.transform.Find("Armature/Hips/Spine/Chest/Neck/SpectrumCollar/Collar").GetComponent<MeshRenderer>().sharedMaterial);

            Texture2D collarNewTex = new Texture2D(2048, 2048);

                collarNewTex = AddTexture(collarNewTex, "_MainTex", Color.white, collarMat, null, 2048);
                collarNewTex = AddTexture(collarNewTex, "_DecalTexture", collarCol, collarMat, null, 2048);

                if (collarFade)
                {
                    collarNewTex = AddTexture(collarNewTex, "_DecalTexture1", collarFadeCol, collarMat, null, 2048);
                }

            collarNewTex = AddTexture(collarNewTex, "_Matcap2Mask", bellCol, collarMat, null, 2048);
            collarNewTex = AddTexture(collarNewTex, "_MatcapMask", metalCol, collarMat, null, 2048);
            collarNewTex = OverlayTexture(collarNewTex, "_DecalTexture3", collarMat, 2048);

            string collarPath = SaveAssets(collarNewTex, questCollarMat, "QuestCollar", preset);

            avatarQ.transform.Find("Armature/Hips/Spine/Chest/Neck/SpectrumCollar/Collar").GetComponent<MeshRenderer>().sharedMaterial = AssetDatabase.LoadAssetAtPath<Material>(collarPath);

        return;
    }

    string SaveAssets(Texture2D texture, Material mat, string name, Preset preset)
    {
        // Encode texture into PNG
        byte[] bytes = texture.EncodeToPNG();

        //generate new material
        string path = AssetDatabase.GenerateUniqueAssetPath("Assets/Spectrum/Generated/" + name + ".mat");
        AssetDatabase.CreateAsset(mat, path);

        //get a path for the texture using the same name as the material
        string texPath = path.Replace("Assets", "").Replace(".mat", ".png");

        //Create PNG texture in generated folder
        File.WriteAllBytes(Application.dataPath + texPath, bytes);

        //refresh files
        AssetDatabase.Refresh();

        preset.ApplyTo(AssetImporter.GetAtPath(path.Replace(".mat", ".png")));

        //set the material's texture to the newly generated texture
        mat.SetTexture("_MainTex", AssetDatabase.LoadAssetAtPath<Texture2D>(path.Replace(".mat", ".png")));

        Debug.Log("Generated " + name + " at " + path);

        return path;
    }

    Texture2D AddTexture(Texture2D texture, string path, Color setColor, Material mat, Material questMat, int resolution)
    {
        Texture2D originalTex = new Texture2D(2048, 2048);
        if (questMat == null && path != "")
        {
            //get the current decal
            originalTex = mat.GetTexture(path) as Texture2D;
        }

        //create a new texture from decal
        Texture2D newTex = new Texture2D(originalTex.width, originalTex.height);
        newTex.SetPixels(originalTex.GetPixels());
        newTex.Apply();
        if (newTex.width != resolution)
        {
            newTex = Resize(newTex, resolution, resolution);
        }
        newTex.Apply();

        //pixels for the selected image/decal
        Color[] newPix = newTex.GetPixels();
        //pixels for the current texture version
        Color[] pix = null;

        if (texture != null)
        {
            pix = texture.GetPixels();
        }

        //color the decal texture
        for (int i = 0; i < newPix.Length; i++)
        {
            Color col = setColor;

            //convert from HDR
            col = new Color(
            Mathf.Pow(col.r, 0.4545f),
            Mathf.Pow(col.g, 0.4545f),
            Mathf.Pow(col.b, 0.4545f),
            255);

            if (path != "")
            {
                col *= newPix[i];
            }

            if (texture != null)
            {
                //if decal, lerp old pixel color and new pixel color by alpha value
                newPix[i] = Color.Lerp(pix[i], col, newPix[i].a);
            }
            else
            {
                newPix[i] = col;
            }

        }

        newTex.SetPixels(newPix);
        newTex.Apply();

        return newTex;
    }

    Texture2D OverlayTexture(Texture2D texture, string path, Material mat, int resolution)
    {
        //get the current decal
        Texture2D originalTex = mat.GetTexture(path) as Texture2D;

        //create a new texture from decal
        Texture2D newTex = new Texture2D(originalTex.width, originalTex.height);
        newTex.SetPixels(originalTex.GetPixels());
        newTex.Apply();
        if (newTex.width != resolution)
        {
            newTex = Resize(newTex, resolution, resolution);
        }
        newTex.Apply();

        //pixels for the selected image/decal
        Color[] newPix = newTex.GetPixels();
        //pixels for the current texture version
        Color[] pix = null;

        if (texture != null)
        {
            pix = texture.GetPixels();
        }

        //overlay the decal texture
        for (int i = 0; i < newPix.Length; i++)
        {
            newPix[i] = Color.Lerp(pix[i], newPix[i], Mathf.Pow(newPix[i].a, 0.4545f));
        }

        newTex.SetPixels(newPix);
        newTex.Apply();

        return newTex;
    }

    //Credit: https://stackoverflow.com/questions/56949217/how-to-resize-a-texture2d-using-height-and-width
    Texture2D Resize(Texture2D texture2D, int targetX, int targetY)
    {
        RenderTexture rt = new RenderTexture(targetX, targetY, 24);
        RenderTexture.active = rt;
        Graphics.Blit(texture2D, rt);
        Texture2D result = new Texture2D(targetX, targetY);
        result.ReadPixels(new Rect(0, 0, targetX, targetY), 0, 0);
        result.Apply();
        return result;
    }
}