namespace Common;

using System.Text;
using Veldrid;
using Veldrid.SPIRV;

    /// <summary>
    /// Contains extension methods for loading <see cref="Shader"/> modules from SPIR-V bytecode.
    /// </summary>
    public static class RFExt
    {
        /// <summary>
        /// Creates a vertex and fragment shader pair from the given <see cref="ShaderDescription"/> pair containing SPIR-V
        /// bytecode or GLSL source code.
        /// </summary>
        /// <param name="factory">The <see cref="ResourceFactory"/> used to compile the translated shader code.</param>
        /// <param name="vertexShaderDescription">The vertex shader's description. <see cref="ShaderDescription.ShaderBytes"/>
        /// should contain SPIR-V bytecode or Vulkan-style GLSL source code which can be compiled to SPIR-V.</param>
        /// <param name="fragmentShaderDescription">The fragment shader's description.
        /// <see cref="ShaderDescription.ShaderBytes"/> should contain SPIR-V bytecode or Vulkan-style GLSL source code which
        /// can be compiled to SPIR-V.</param>
        /// <returns>A two-element array, containing the vertex shader (element 0) and the fragment shader (element 1).</returns>
        public static Shader[] CrossCompile(
            this ResourceFactory factory,
            ShaderDescription vertexShaderDescription,
            ShaderDescription fragmentShaderDescription)
        {
            return CrossCompile(factory, vertexShaderDescription, fragmentShaderDescription, new CrossCompileOptions());
        }

        /// <summary>
        /// Creates a vertex and fragment shader pair from the given <see cref="ShaderDescription"/> pair containing SPIR-V
        /// bytecode or GLSL source code.
        /// </summary>
        /// <param name="factory">The <see cref="ResourceFactory"/> used to compile the translated shader code.</param>
        /// <param name="vertexShaderDescription">The vertex shader's description. <see cref="ShaderDescription.ShaderBytes"/>
        /// should contain SPIR-V bytecode or Vulkan-style GLSL source code which can be compiled to SPIR-V.</param>
        /// <param name="fragmentShaderDescription">The fragment shader's description.
        /// <see cref="ShaderDescription.ShaderBytes"/> should contain SPIR-V bytecode or Vulkan-style GLSL source code which
        /// can be compiled to SPIR-V.</param>
        /// <param name="options">The <see cref="CrossCompileOptions"/> which will control the parameters used to translate the
        /// shaders from SPIR-V to the target language.</param>
        /// <returns>A two-element array, containing the vertex shader (element 0) and the fragment shader (element 1).</returns>
        public static Shader[] CrossCompile(
            this ResourceFactory factory,
            ShaderDescription vertexShaderDescription,
            ShaderDescription fragmentShaderDescription,
            CrossCompileOptions options)
        {
            options.NormalizeResourceNames = true;
            GraphicsBackend backend = factory.BackendType;
            if (backend == GraphicsBackend.Vulkan)
            {
                vertexShaderDescription.ShaderBytes = EnsureSpirv(backend, vertexShaderDescription);
                fragmentShaderDescription.ShaderBytes = EnsureSpirv(backend, fragmentShaderDescription);

                return new Shader[]
                {
                    factory.CreateShader(ref vertexShaderDescription),
                    factory.CreateShader(ref fragmentShaderDescription)
                };
            }

            CrossCompileTarget target = GetCompilationTarget(factory.BackendType);
            VertexFragmentCompilationResult compilationResult = SpirvCompilation.CompileVertexFragment(
                vertexShaderDescription.ShaderBytes,
                fragmentShaderDescription.ShaderBytes,
                target,
                options);

            string vertexEntryPoint = (backend == GraphicsBackend.Metal && vertexShaderDescription.EntryPoint == "main")
                ? "main0"
                : vertexShaderDescription.EntryPoint;
            byte[] vertexBytes = GetBytes(backend, compilationResult.VertexShader);
            Shader vertexShader = factory.CreateShader(new ShaderDescription(
                vertexShaderDescription.Stage,
                vertexBytes,
                vertexEntryPoint)
                {
                    Debug = vertexShaderDescription.Debug
                });

            string fragmentEntryPoint = (backend == GraphicsBackend.Metal && fragmentShaderDescription.EntryPoint == "main")
                ? "main0"
                : fragmentShaderDescription.EntryPoint;
            byte[] fragmentBytes = GetBytes(backend, compilationResult.FragmentShader);
            Shader fragmentShader = factory.CreateShader(new ShaderDescription(
                fragmentShaderDescription.Stage,
                fragmentBytes,
                fragmentEntryPoint)
                {
                    Debug = fragmentShaderDescription.Debug
                });

            return new Shader[] { vertexShader, fragmentShader };
        }

        /// <summary>
        /// Creates a compute shader from the given <see cref="ShaderDescription"/> containing SPIR-V bytecode or GLSL source
        /// code.
        /// </summary>
        /// <param name="factory">The <see cref="ResourceFactory"/> used to compile the translated shader code.</param>
        /// <param name="computeShaderDescription">The compute shader's description.
        /// <see cref="ShaderDescription.ShaderBytes"/> should contain SPIR-V bytecode or Vulkan-style GLSL source code which
        /// can be compiled to SPIR-V.</param>
        /// <returns>The compiled compute <see cref="Shader"/>.</returns>
        public static Shader CrossCompile(
            this ResourceFactory factory,
            ShaderDescription computeShaderDescription)
        {
            return CrossCompile(factory, computeShaderDescription, new CrossCompileOptions());
        }

        /// <summary>
        /// Creates a compute shader from the given <see cref="ShaderDescription"/> containing SPIR-V bytecode or GLSL source
        /// code.
        /// </summary>
        /// <param name="factory">The <see cref="ResourceFactory"/> used to compile the translated shader code.</param>
        /// <param name="computeShaderDescription">The compute shader's description.
        /// <see cref="ShaderDescription.ShaderBytes"/> should contain SPIR-V bytecode or Vulkan-style GLSL source code which
        /// can be compiled to SPIR-V.</param>
        /// <param name="options">The <see cref="CrossCompileOptions"/> which will control the parameters used to translate the
        /// shaders from SPIR-V to the target language.</param>
        /// <returns>The compiled compute <see cref="Shader"/>.</returns>
        public static Shader CrossCompile(
            this ResourceFactory factory,
            ShaderDescription computeShaderDescription,
            CrossCompileOptions options)
        {
            GraphicsBackend backend = factory.BackendType;
            if (backend == GraphicsBackend.Vulkan)
            {
                computeShaderDescription.ShaderBytes = EnsureSpirv(backend, computeShaderDescription);
                return factory.CreateShader(ref computeShaderDescription);
            }

            CrossCompileTarget target = GetCompilationTarget(factory.BackendType);
            ComputeCompilationResult compilationResult = SpirvCompilation.CompileCompute(
                computeShaderDescription.ShaderBytes,
                target,
                options);

            string computeEntryPoint = (backend == GraphicsBackend.Metal && computeShaderDescription.EntryPoint == "main")
                ? "main0"
                : computeShaderDescription.EntryPoint;
            byte[] computeBytes = GetBytes(backend, compilationResult.ComputeShader);
            return factory.CreateShader(new ShaderDescription(
                computeShaderDescription.Stage,
                computeBytes,
                computeEntryPoint)
                {
                    Debug = computeShaderDescription.Debug
                });
        }

        private static unsafe byte[] EnsureSpirv(GraphicsBackend backend, ShaderDescription description)
        {
            if (HasSpirvHeader(description.ShaderBytes))
            {
                return description.ShaderBytes;
            }
            else
            {
                var src = GetSource(backend, description.ShaderBytes);
                var result = SpirvCompilation.CompileGlslToSpirv(
                    src,
                    null,
                    description.Stage,
                    new GlslCompileOptions(description.Debug)
                );
                return result.SpirvBytes;
            }
        }

        private static bool HasSpirvHeader(byte[] bytes)
        {
            return bytes.Length > 4
                && bytes[0] == 0x03
                && bytes[1] == 0x02
                && bytes[2] == 0x23
                && bytes[3] == 0x07;
        }

        private static byte[] GetBytes(GraphicsBackend backend, string code)
        {
            switch (backend)
            {
                case GraphicsBackend.Direct3D11:
                case GraphicsBackend.OpenGL:
                case GraphicsBackend.OpenGLES:
                    return Encoding.ASCII.GetBytes(code);
                case GraphicsBackend.Metal:
                    return Encoding.UTF8.GetBytes(code);
                default:
                    throw new SpirvCompilationException($"Invalid GraphicsBackend: {backend}");
            }
        }

        private static string GetSource(GraphicsBackend backend, byte[] code)
        {
            switch (backend)
            {
                case GraphicsBackend.Direct3D11:
                case GraphicsBackend.OpenGL:
                case GraphicsBackend.OpenGLES:
                    return Encoding.ASCII.GetString(code);
                case GraphicsBackend.Metal:
                    return Encoding.UTF8.GetString(code);
                default:
                    throw new SpirvCompilationException($"Invalid GraphicsBackend: {backend}");
            }
        }

        private static CrossCompileTarget GetCompilationTarget(GraphicsBackend backend)
        {
            switch (backend)
            {
                case GraphicsBackend.Direct3D11:
                    return CrossCompileTarget.HLSL;
                case GraphicsBackend.OpenGL:
                    return CrossCompileTarget.GLSL;
                case GraphicsBackend.Metal:
                    return CrossCompileTarget.MSL;
                case GraphicsBackend.OpenGLES:
                    return CrossCompileTarget.ESSL;
                default:
                    throw new SpirvCompilationException($"Invalid GraphicsBackend: {backend}");
            }
        }
    }