using System.IO;
using UnityEditor;
using UnityEngine;
using UnityEngine.Experimental.Rendering;
using UnityEngine.Rendering;
using UnityEngine.UIElements;

namespace Misaki.HdrpToon.Editor
{
    internal class SSSLutBakerView : UnityEditor.EditorWindow
    {
        public enum Quality
        {
            [InspectorName("64")]
            UltraLow = 64,
            [InspectorName("128")]
            VeryLow = 128,
            [InspectorName("256")]
            Low = 256,
            [InspectorName("512")]
            Medium = 512,
            [InspectorName("1024")]
            High = 1024,
            [InspectorName("2048")]
            VeryHigh = 2048,
            [InspectorName("4096")]
            UltraHigh = 4096,
        }

        private const int _PREVIEW_SIZE = 128;

        [SerializeField]
        private VisualTreeAsset _visualAsset;
        [SerializeField]
        private ComputeShader _bakerShader;

        private RenderTexture _previewTexture;
        private RenderTextureDescriptor _previewDescriptor;

        public bool isShadowLut;
        public float scatterRadius = 1.0f;
        public float positionShift = 0.0f;
        public float intensity = 1.0f;
        public Quality sampleCount = Quality.Low;
        public bool applyTonemap = true;
        public Quality outputResolution = Quality.Medium;
        public int outputQuality = 100;

        [MenuItem("Tools/UTS/SSS LUT Baker")]
        private static void ShowWindow()
        {
            var window = GetWindow<SSSLutBakerView>(true, "SSS LUT Baker");
            window.minSize = new Vector2(400, 650);
            window.ShowUtility();
        }

        private void OnEnable()
        {
            _previewDescriptor = new RenderTextureDescriptor
            {
                width = _PREVIEW_SIZE,
                height = _PREVIEW_SIZE,
                volumeDepth = 1,
                dimension = TextureDimension.Tex2D,
                depthBufferBits = 0,
                msaaSamples = 1,
                graphicsFormat = GraphicsFormat.R8G8B8A8_UNorm,
                enableRandomWrite = true,
                sRGB = false,
                useMipMap = false,
            };

            _previewTexture = RenderTexture.GetTemporary(_previewDescriptor);
        }

        private bool GenerateSSSLut(RenderTexture texture, int textureSize, bool production)
        {
            if (_bakerShader == null || !texture.enableRandomWrite)
            {
                return false;
            }

            var kernelIndex = isShadowLut ? 1 : 0;

            _bakerShader.SetFloat("_ScatterRadius", scatterRadius);
            _bakerShader.SetFloat("_PositionShift", positionShift);
            _bakerShader.SetFloat("_Intensity", intensity);
            _bakerShader.SetInt("_SampleCount", (int)sampleCount);
            _bakerShader.SetInt("_ApplyTonemap", applyTonemap ? 1 : 0);
            _bakerShader.SetInt("_TextureSize", textureSize);

            _bakerShader.SetKeyword(new LocalKeyword(_bakerShader, "_PRODUCTION"), production);

            _bakerShader.SetTexture(kernelIndex, "_SSSLut", texture);

            const int groupSizeX = 8;
            const int groupSizeY = 8;
            var threadGroupX = (textureSize + (groupSizeX - 1)) / groupSizeX;
            var threadGroupY = (textureSize + (groupSizeY - 1)) / groupSizeY;
            _bakerShader.Dispatch(kernelIndex, threadGroupX, threadGroupY, 1);

            return true;
        }

        private void BakePreview()
        {
            GenerateSSSLut(_previewTexture, _PREVIEW_SIZE, false);
        }

        private void Bake()
        {
            var outputPath = EditorUtility.SaveFilePanel("Export Texture", Application.dataPath, "SSS Lut", "jpg");
            if (string.IsNullOrEmpty(outputPath))
            {
                return;
            }

            var outputDescriptor = _previewDescriptor;
            outputDescriptor.width = (int)outputResolution;
            outputDescriptor.height = (int)outputResolution;

            var tempTexture = RenderTexture.GetTemporary(outputDescriptor);
            var tempTexture2D = new Texture2D(outputDescriptor.width, outputDescriptor.height, TextureFormat.RGB24, false, true);

            try
            {
                if (!GenerateSSSLut(tempTexture, (int)outputResolution, true))
                {
                    return;
                }

                RenderTexture.active = tempTexture;
                tempTexture2D.ReadPixels(new Rect(0, 0, outputDescriptor.width, outputDescriptor.height), 0, 0);
                tempTexture2D.Apply();
                RenderTexture.active = null;

                File.WriteAllBytes(outputPath, tempTexture2D.EncodeToJPG(outputQuality));
                AssetDatabase.Refresh();
            }
            finally
            {
                RenderTexture.ReleaseTemporary(tempTexture);
                DestroyImmediate(tempTexture2D);
            }
        }

        private void CreateGUI()
        {
            if (_visualAsset == null)
            {
                return;
            }

            var visualTree = _visualAsset.Instantiate();
            visualTree.StretchToParentSize();
            visualTree.dataSource = this;

            var previewImage = visualTree.Q<Image>("preview-image");
            previewImage.image = _previewTexture;

            var previewButton = visualTree.Q<Button>("preview-button");
            previewButton.clickable.clicked += BakePreview;

            var bakeButton = visualTree.Q<Button>("bake-button");
            bakeButton.clickable.clicked += Bake;

            rootVisualElement.Add(visualTree);
        }

        private void OnDestroy()
        {
            RenderTexture.ReleaseTemporary(_previewTexture);
        }
    }
}