using Vortice.ShaderCompiler;
using Vortice.SPIRV.Reflect;

namespace Nerfed.Builder;

// Values should match the ShaderStage enum in the runtime.
public enum ShaderStage
{
    Vertex,
    Fragment
}

// Values should match the ShaderFormat enum in the runtime.
public enum ShaderFormat
{
    Invalid,
    SPIRV,
    HLSL,
    DXBC,
    DXIL,
    MSL,
    METALLIB,
    SECRET
}

public class ShaderImporter : IImporter
{
    private readonly ShaderStage shaderStage;

    public ShaderImporter(ShaderStage shaderStage)
    {
        this.shaderStage = shaderStage;
    }

    public unsafe void Import(string inFile, string outFile)
    {
        string name = Path.GetFileNameWithoutExtension(inFile);
        string nameWithExt = Path.GetFileName(inFile);

        // Compile the shader.
        Result compileResult;
        using (Compiler compiler = new Compiler())
        {
            string shaderSource = File.ReadAllText(inFile);
            compileResult = compiler.Compile(shaderSource, nameWithExt, ToShaderKind(shaderStage));
        }

        if (compileResult.Status != CompilationStatus.Success)
        {
            Console.Error.WriteLine($"Failed to compile {nameWithExt}\n{compileResult.ErrorMessage}");
            return;
        }

        if (compileResult.ErrorsCount > 0 || compileResult.WarningsCount > 0)
        {
            Console.Error.WriteLine(compileResult.ErrorMessage);
        }

        Span<byte> byteCode = compileResult.GetBytecode();

        // Inspect SPIR-V bytecode for information which the runtime requires to create a shader resource.
        SpvReflectShaderModule module = new SpvReflectShaderModule();
        if (!CheckReflectResult(SPIRVReflectApi.spvReflectCreateShaderModule(byteCode, &module), name))
        {
            return;
        }

        uint descriptorSetCount = 0;
        if (!CheckReflectResult(SPIRVReflectApi.spvReflectEnumerateDescriptorSets(&module, &descriptorSetCount, null), name))
        {
            return;
        }

        int uniformBufferCount = 0;
        int storageBufferCount = 0;
        int storageTextureCount = 0;
        int samplerCount = 0;

        if (descriptorSetCount > 0)
        {
            SpvReflectDescriptorSet* descriptorSets = stackalloc SpvReflectDescriptorSet[(int)descriptorSetCount];
            if (!CheckReflectResult(
                    SPIRVReflectApi.spvReflectEnumerateDescriptorSets(&module, &descriptorSetCount, &descriptorSets),
                    name
                ))
            {
                return;
            }

            for (int i = 0; i < descriptorSetCount; i++)
            {
                SpvReflectDescriptorSet set = descriptorSets[i];
                for (int j = 0; j < set.binding_count; j++)
                {
                    SpvReflectDescriptorBinding binding = *set.bindings[j];
                    if (binding.descriptor_type == SpvReflectDescriptorType.UniformBuffer)
                    {
                        uniformBufferCount++;
                    }
                    else if (binding.descriptor_type == SpvReflectDescriptorType.StorageBuffer)
                    {
                        storageBufferCount++;
                    }
                    else if (binding.descriptor_type == SpvReflectDescriptorType.StorageTexelBuffer)
                    {
                        storageTextureCount++;
                    }
                    else if (binding.descriptor_type == SpvReflectDescriptorType.Sampler ||
                             binding.descriptor_type == SpvReflectDescriptorType.CombinedImageSampler)
                    {
                        samplerCount++;
                    }
                }
            }
        }

        //TODO: Convert SPIR-V to other bytecode formats here (DX/Consoles).
        ShaderFormat format = ShaderFormat.SPIRV;

        // Write shader meta-data and bytecode to the output file.
        using (FileStream stream = new FileStream(outFile, FileMode.Create, FileAccess.Write))
        {
            using (BinaryWriter writer = new BinaryWriter(stream))
            {
                writer.Write((int)format);
                writer.Write((int)shaderStage);
                writer.Write(uniformBufferCount);
                writer.Write(storageBufferCount);
                writer.Write(storageTextureCount);
                writer.Write(samplerCount);
                writer.Write(byteCode.Length);
                writer.Write(byteCode);
            }
        }
    }

    private bool CheckReflectResult(SpvReflectResult result, string name)
    {
        if (result != SpvReflectResult.Success)
        {
            Console.Error.WriteLine($"SpirV-Reflect failure for '{name}': {result}");
            return false;
        }

        return true;
    }

    private ShaderKind ToShaderKind(ShaderStage shaderStage)
    {
        switch (shaderStage)
        {
            case ShaderStage.Vertex: return ShaderKind.VertexShader;
            case ShaderStage.Fragment: return ShaderKind.FragmentShader;
            default: throw new ArgumentOutOfRangeException(nameof(shaderStage));
        }
    }
}