prism/engine/gfx/webgpu/src/gfx_webgpu.cpp

#include "gfx_webgpu.hpp"

#include "gfx_commandbuffer.hpp"
#include "gfx_webgpu_buffer.hpp"
#include "gfx_webgpu_commandbuffer.hpp"
#include "gfx_webgpu_framebuffer.hpp"
#include "gfx_webgpu_pipeline.hpp"
#include "gfx_webgpu_renderpass.hpp"
#include "gfx_webgpu_sampler.hpp"
#include "gfx_webgpu_texture.hpp"
#include "utility.hpp"

WGPUTextureFormat toPixFormat(GFXPixelFormat format) {
    switch (format) {
        case GFXPixelFormat::R_32F:
            return WGPUTextureFormat_R32Float;

        case GFXPixelFormat::R_16F:
            return WGPUTextureFormat_R16Float;

        case GFXPixelFormat::RGBA_32F:
            return WGPUTextureFormat_RGBA32Float;

        case GFXPixelFormat::RGBA8_UNORM:
            return WGPUTextureFormat_RGBA8Unorm;

        case GFXPixelFormat::R8_UNORM:
            return WGPUTextureFormat_R8Unorm;

        case GFXPixelFormat::R8G8_UNORM:
            return WGPUTextureFormat_RG8Unorm;

        case GFXPixelFormat::R8G8_SFLOAT:
            return WGPUTextureFormat_RG16Float;

        case GFXPixelFormat::R8G8B8A8_UNORM:
            return WGPUTextureFormat_RGBA8Unorm;

        case GFXPixelFormat::R16G16B16A16_SFLOAT:
            return WGPUTextureFormat_RGBA16Float;

        case GFXPixelFormat::DEPTH_32F:
            return WGPUTextureFormat_Depth32Float;
    }

    return WGPUTextureFormat_Undefined;
}

WGPUVertexFormat toVertFormat(GFXVertexFormat format) {
    switch (format) {
        case GFXVertexFormat::FLOAT2:
            return WGPUVertexFormat_Float32x2;
        case GFXVertexFormat::FLOAT3:
            return WGPUVertexFormat_Float32x3;
        case GFXVertexFormat::FLOAT4:
            return WGPUVertexFormat_Float32x4;
        case GFXVertexFormat::INT:
            return WGPUVertexFormat_Sint32;
        case GFXVertexFormat::INT4:
            return WGPUVertexFormat_Sint32x4;
        case GFXVertexFormat::UNORM4:
            return WGPUVertexFormat_Unorm8x4;
    }

    return WGPUVertexFormat_Undefined;
}

WGPUBlendFactor toFactor(GFXBlendFactor factor) {
    switch (factor) {
        case GFXBlendFactor::Zero:
            return WGPUBlendFactor_Zero;
        case GFXBlendFactor::One:
            return WGPUBlendFactor_One;
        case GFXBlendFactor::OneMinusSrcAlpha:
            return WGPUBlendFactor_OneMinusSrcAlpha;
        case GFXBlendFactor::OneMinusSrcColor:
            return WGPUBlendFactor_OneMinusSrc;
        case GFXBlendFactor::SrcAlpha:
            return WGPUBlendFactor_SrcAlpha;
        case GFXBlendFactor::DstAlpha:
            return WGPUBlendFactor_DstAlpha;
        case GFXBlendFactor::SrcColor:
            return WGPUBlendFactor_Src;
        case GFXBlendFactor::DstColor:
            return WGPUBlendFactor_Dst;
    }

    return WGPUBlendFactor_One;
}

WGPUAddressMode toSamplerMode(SamplingMode mode) {
    switch (mode) {
        case SamplingMode::Repeat:
            return WGPUAddressMode_Repeat;
        case SamplingMode::ClampToEdge:
            return WGPUAddressMode_ClampToEdge;
    }

    return WGPUAddressMode_Repeat;
}

WGPUFilterMode toFilter(GFXFilter filter) {
    switch (filter) {
        case GFXFilter::Nearest:
            return WGPUFilterMode_Nearest;
        case GFXFilter::Linear:
            return WGPUFilterMode_Linear;
    }

    return WGPUFilterMode_Linear;
}

WGPUCompareFunction toCompareFunc(GFXCompareFunction func) {
    switch (func) {
        case GFXCompareFunction::Never:
            return WGPUCompareFunction_Never;
        case GFXCompareFunction::Less:
            return WGPUCompareFunction_Less;
        case GFXCompareFunction::Equal:
            return WGPUCompareFunction_Equal;
        case GFXCompareFunction::LessOrEqual:
            return WGPUCompareFunction_LessEqual;
        case GFXCompareFunction::Greater:
            return WGPUCompareFunction_Greater;
        case GFXCompareFunction::NotEqual:
            return WGPUCompareFunction_NotEqual;
        case GFXCompareFunction::GreaterOrEqual:
            return WGPUCompareFunction_GreaterEqual;
        case GFXCompareFunction::Always:
            return WGPUCompareFunction_Always;
    }
}

WGPUShaderModule GFXWebGPU::create_shader(const uint32_t* code, const uint32_t size, std::string_view label) {
    WGPUShaderModuleSPIRVDescriptor spirv = {};
    spirv.chain.sType = WGPUSType_ShaderModuleSPIRVDescriptor;
    spirv.codeSize = size / sizeof(uint32_t);
    spirv.code = code;

    WGPUShaderModuleDescriptor desc = {};
    desc.nextInChain = reinterpret_cast<WGPUChainedStruct*>(&spirv);
    desc.label = label.data();

    return wgpuDeviceCreateShaderModule(device, &desc);
}

WGPUSwapChain GFXWebGPU::create_swapchain() {
    WGPUSurfaceDescriptorFromCanvasHTMLSelector canvDesc = {};
    canvDesc.chain.sType = WGPUSType_SurfaceDescriptorFromCanvasHTMLSelector;
    canvDesc.selector = "canvas";

    WGPUSurfaceDescriptor surfDesc = {};
    surfDesc.nextInChain = reinterpret_cast<WGPUChainedStruct*>(&canvDesc);

    WGPUSurface surface = wgpuInstanceCreateSurface(nullptr, &surfDesc);

    WGPUSwapChainDescriptor swapDesc = {};
    swapDesc.usage  = WGPUTextureUsage_RenderAttachment;
    swapDesc.format = WGPUTextureFormat_BGRA8Unorm;
    swapDesc.width  = 800; // TODO: configurable
    swapDesc.height = 450;
    swapDesc.presentMode = WGPUPresentMode_Fifo;

    return wgpuDeviceCreateSwapChain(device, surface, &swapDesc);
}

bool GFXWebGPU::initialize(const GFXCreateInfo& createInfo) {
    device = emscripten_webgpu_get_device();
    if(device == nullptr) {
        prism::log("Failed to get WebGPU device!");
    }

    queue = wgpuDeviceGetQueue(device);
    swapchain = create_swapchain();

    prism::log("Initialized WebGPU!");

    return true;
}

ShaderLanguage GFXWebGPU::accepted_shader_language() {
    return ShaderLanguage::WGSL;
}

const char* GFXWebGPU::get_name() {
    return "WebGPU";
}

GFXCommandBuffer* presentation_command_buffer = nullptr;

GFXCommandBuffer* GFXWebGPU::acquire_command_buffer(bool for_presentation_use) {
    if(for_presentation_use) {
        if(presentation_command_buffer == nullptr)
            presentation_command_buffer = new GFXCommandBuffer();

        presentation_command_buffer->commands.clear();

        return presentation_command_buffer;
    } else {
        return new GFXCommandBuffer();
    }
}

GFXBuffer* GFXWebGPU::create_buffer(void* data, const GFXSize size, const GFXBufferUsage usage) {
    auto buffer = new GFXWebGPUBuffer();
    buffer->size = size;

    WGPUBufferDescriptor desc = {};
    desc.size = size;
    desc.mappedAtCreation = true;

    switch(usage) {
        case GFXBufferUsage::Vertex:
            desc.usage = WGPUBufferUsage_Vertex;
            break;
        case GFXBufferUsage::Index:
            desc.usage = WGPUBufferUsage_Index;
            break;
        case GFXBufferUsage::Storage:
            desc.usage = WGPUBufferUsage_Storage;
            break;
    }

    desc.usage |= WGPUBufferUsage_Uniform; // TODO: not currently exposed in gfx api
    desc.usage |= WGPUBufferUsage_CopyDst;

    buffer->handle = wgpuDeviceCreateBuffer(device, &desc);

    if(data != nullptr) {
        wgpuQueueWriteBuffer(queue, buffer->handle, 0, data, size);
    }

    return buffer;
}

void GFXWebGPU::copy_buffer(GFXBuffer* buffer, void* data, const GFXSize offset, const GFXSize size) {
    auto wgpu_buffer = (GFXWebGPUBuffer*)buffer;

    wgpuQueueWriteBuffer(queue, wgpu_buffer->handle, offset, data, size);
}

void* GFXWebGPU::get_buffer_contents(GFXBuffer* buffer) {
    auto wgpu_buffer = (GFXWebGPUBuffer*)buffer;

    return wgpuBufferGetMappedRange(wgpu_buffer->handle, 0, wgpu_buffer->size);
}

void GFXWebGPU::release_buffer_contents(GFXBuffer* buffer, void* handle) {
    auto wgpu_buffer = (GFXWebGPUBuffer*)buffer;
}

GFXTexture* GFXWebGPU::create_texture(const GFXTextureCreateInfo& info) {
    auto texture = new GFXWebGPUTexture();

    int array_length = info.array_length;
    if (info.type == GFXTextureType::Cubemap)
        array_length = 6;
    else if (info.type == GFXTextureType::CubemapArray)
        array_length *= 6;

    WGPUTextureDescriptor descriptor = {};
    descriptor.size.width = info.width;
    descriptor.size.height = info.height;
    descriptor.size.depthOrArrayLayers = array_length;

    descriptor.dimension = WGPUTextureDimension_2D;

    descriptor.format = WGPUTextureFormat_BGRA8Unorm;

    descriptor.label = info.label.c_str();

    descriptor.mipLevelCount = info.mip_count;

    descriptor.usage = WGPUTextureUsage_None;
    descriptor.sampleCount = 1;

    if((info.usage & GFXTextureUsage::Sampled) == GFXTextureUsage::Sampled)
        descriptor.usage |= WGPUTextureUsage_TextureBinding;

    if((info.usage & GFXTextureUsage::Attachment) == GFXTextureUsage::Attachment)
        descriptor.usage |= WGPUTextureUsage_RenderAttachment;

    if((info.usage & GFXTextureUsage::TransferDst) == GFXTextureUsage::TransferDst)
        descriptor.usage |= WGPUTextureUsage_CopyDst;

    if((info.usage & GFXTextureUsage::TransferSrc) == GFXTextureUsage::TransferSrc)
        descriptor.usage |= WGPUTextureUsage_CopySrc;

    if((info.usage & GFXTextureUsage::Storage) == GFXTextureUsage::Storage)
        descriptor.usage |= WGPUTextureUsage_StorageBinding;

    texture->handle = wgpuDeviceCreateTexture(device, &descriptor);

    WGPUTextureViewDescriptor view_descriptor = {};
    view_descriptor.label = info.label.data();
    view_descriptor.format = descriptor.format;
    view_descriptor.mipLevelCount = descriptor.mipLevelCount;
    view_descriptor.baseArrayLayer = 0;
    view_descriptor.aspect = WGPUTextureAspect_All;
    view_descriptor.arrayLayerCount = array_length;

    switch(info.type) {
        case GFXTextureType::Single2D:
            view_descriptor.dimension = WGPUTextureViewDimension_2D;
            break;
        case GFXTextureType::Array2D:
            view_descriptor.dimension = WGPUTextureViewDimension_2DArray;
            break;
        case GFXTextureType::Cubemap:
            view_descriptor.dimension = WGPUTextureViewDimension_Cube;
            break;
        case GFXTextureType::CubemapArray:
            view_descriptor.dimension = WGPUTextureViewDimension_CubeArray;
            break;
    }

    texture->view = wgpuTextureCreateView(texture->handle, &view_descriptor);

    return texture;
}

void GFXWebGPU::copy_texture(GFXTexture* texture, void* data, GFXSize size) {
    GFX::copy_texture(texture, data, size);
}

void GFXWebGPU::copy_texture(GFXTexture* from, GFXTexture* to) {
    GFX::copy_texture(from, to);
}

void GFXWebGPU::copy_texture(GFXTexture* from, GFXBuffer* to) {
    GFX::copy_texture(from, to);
}

GFXSampler* GFXWebGPU::create_sampler(const GFXSamplerCreateInfo& info) {
    auto sampler = new GFXWebGPUSampler();

    WGPUSamplerDescriptor sampler_descriptor = {};

    sampler->handle = wgpuDeviceCreateSampler(device, &sampler_descriptor);

    return sampler;
}

GFXFramebuffer* GFXWebGPU::create_framebuffer(const GFXFramebufferCreateInfo& info) {
    auto framebuffer = new GFXWebGPUFramebuffer();

    for(auto attachment : info.attachments) {
        framebuffer->attachments.push_back((GFXWebGPUTexture*)attachment);
    }

    return framebuffer;
}

GFXRenderPass* GFXWebGPU::create_render_pass(const GFXRenderPassCreateInfo& info) {
    auto render_pass = new GFXWebGPURenderPass();

    render_pass->attachments = info.attachments;

    return render_pass;
}

GFXPipeline* GFXWebGPU::create_graphics_pipeline(const GFXGraphicsPipelineCreateInfo& info) {
    auto pipeline = new GFXWebGPUPipeline();

    pipeline->label = info.label;

    WGPURenderPipelineDescriptor descriptor = {};
    descriptor.label = pipeline->label.c_str();

    const bool has_vertex_stage = !info.shaders.vertex_src.empty();
    if (has_vertex_stage) {
        const bool vertex_use_shader_source = !info.shaders.vertex_src.is_path();

        if (vertex_use_shader_source) {
            auto vertex_shader_vector = info.shaders.vertex_src.as_bytecode();

            descriptor.vertex.module = create_shader(vertex_shader_vector.data(), vertex_shader_vector.size() * sizeof(uint32_t), std::string(info.label + " vertex stage").c_str());
        }
        else {
            auto vertex_shader = prism::open_file(prism::internal_domain / (info.shaders.vertex_src.as_path().string() + ".wgsl"), true);
            vertex_shader->read_all();

            descriptor.vertex.module = create_shader(vertex_shader->cast_data<uint32_t>(), vertex_shader->size(), info.shaders.vertex_src.as_path().string().c_str());
        }
    }

    std::vector<WGPUColorTargetState> color_targets;

    WGPUFragmentState fragment = {};
    fragment.entryPoint = "main";

    const bool has_fragment_stage = !info.shaders.fragment_src.empty();
    if (has_fragment_stage) {
        descriptor.fragment = &fragment;

        if(info.render_pass == nullptr) {
            // swapchain image
            WGPUColorTargetState target_state = {};
            target_state.format = WGPUTextureFormat_BGRA8Unorm;
            target_state.writeMask = WGPUColorWriteMask_All;

            color_targets.push_back(target_state);

            prism::log("is using swapchain image :-)");
        } else {
            for(auto attachment : ((GFXWebGPURenderPass*)info.render_pass)->attachments) {
                if(attachment != GFXPixelFormat::DEPTH_32F) {
                    WGPUColorTargetState target_state = {};
                    target_state.format = toPixFormat(attachment);
                    target_state.writeMask = WGPUColorWriteMask_All;

                    color_targets.push_back(target_state);
                }
            }
        }

        const bool fragment_use_shader_source = !info.shaders.fragment_src.is_path();

        if (fragment_use_shader_source) {
            auto fragment_shader_vector = info.shaders.fragment_src.as_bytecode();

            fragment.module = create_shader(fragment_shader_vector.data(), fragment_shader_vector.size() * sizeof(uint32_t), std::string(info.label + " fragment stage").c_str());
        }
        else {
            auto fragment_shader = prism::open_file(prism::internal_domain / (info.shaders.fragment_src.as_path().string() + ".wgsl"), true);
            fragment_shader->read_all();

            fragment.module = create_shader(fragment_shader->cast_data<uint32_t>(), fragment_shader->size(), info.shaders.fragment_src.as_path().string().c_str());
        }
    }

    fragment.targetCount = color_targets.size();
    fragment.targets = color_targets.data();

    descriptor.vertex.entryPoint = "main";

    prism::log("building pipeline {}", info.label);
    prism::log("--------");

    // alright, webgpu in their infinite wisdom does not allow arbitrary binding locations
    // so, to be consistent with what vulkan, metal and virtually every other API allows,
    // we will create dummy buffers with no attributes attached. congrats webgpu devs.
    int dummy_buffer_count = 0;
    for (auto& binding : info.vertex_input.inputs) {
        dummy_buffer_count = std::max(binding.location, dummy_buffer_count);
    }
    dummy_buffer_count += 1;

    std::vector<std::vector<WGPUVertexAttribute>> attributes;
    attributes.resize(dummy_buffer_count);

    prism::log("dummy buffer count: {}", dummy_buffer_count);

    for (auto& attribute : info.vertex_input.attributes) {
        WGPUVertexAttribute description;
        description.shaderLocation = attribute.location;
        description.format = toVertFormat(attribute.format);
        description.offset = attribute.offset;

        attributes[attribute.binding].push_back(description);
    }

    std::vector<WGPUVertexBufferLayout> inputs;
    inputs.resize(dummy_buffer_count);
    for (auto& binding : info.vertex_input.inputs) {
        prism::log("binding loc {}", binding.location);

        WGPUVertexBufferLayout b;
        b.attributes = attributes[binding.location].data();
        b.attributeCount = attributes[binding.location].size();
        b.arrayStride = binding.stride;
        b.stepMode = WGPUVertexStepMode_Vertex;

        inputs[binding.location] = b;
    }

    descriptor.vertex.buffers = inputs.data();
    descriptor.vertex.bufferCount = inputs.size();

    switch (info.rasterization.culling_mode) {
        case GFXCullingMode::Backface:
            descriptor.primitive.cullMode = WGPUCullMode_Back;
            break;
        case GFXCullingMode::Frontface:
            descriptor.primitive.cullMode = WGPUCullMode_Front;
            break;
        case GFXCullingMode::None:
            descriptor.primitive.cullMode = WGPUCullMode_None;
    }

    switch (info.rasterization.winding_mode) {
        case GFXWindingMode::Clockwise:
            descriptor.primitive.frontFace = WGPUFrontFace_CW;
            break;
        case GFXWindingMode::CounterClockwise:
            descriptor.primitive.frontFace = WGPUFrontFace_CCW;
            break;
    }

    descriptor.primitive.topology = WGPUPrimitiveTopology_TriangleList;

    // create bind group layout
    std::vector<WGPUBindGroupLayoutEntry> group_entries = {};
    for (auto& binding : info.shader_input.bindings) {
        WGPUBindGroupLayoutEntry entry = {};
        entry.binding = binding.binding;
        entry.visibility = WGPUShaderStage_Vertex | WGPUShaderStage_Fragment;

        switch (binding.type) {
        case GFXBindingType::PushConstant:
        {
            prism::log("Note: enabling push constant emulation for pipeline {}", info.label);

            pipeline->push_constant_emulation_buffer = (GFXWebGPUBuffer*)create_buffer(nullptr,
                                                                                       pipeline->push_constant_emulation_size,
                                                                                       GFXBufferUsage::Storage);

            pipeline->push_constant_emulation_buffer_index = binding.binding;

            entry.buffer.type = WGPUBufferBindingType_Uniform;
        }
            break;
		case GFXBindingType::StorageBuffer:
        {
            entry.buffer.type = WGPUBufferBindingType_Uniform;
        }
			break;
		case GFXBindingType::StorageImage:
		{
            entry.storageTexture.access = WGPUStorageTextureAccess_WriteOnly;
		}
			break;
		case GFXBindingType::SampledImage:
		{
            entry.texture.viewDimension = WGPUTextureViewDimension_2D;
		}
		    break;
		case GFXBindingType::Sampler: {
            entry.sampler.type = WGPUSamplerBindingType_Filtering;
        }
			break;
		}

        group_entries.push_back(entry);
    }

    descriptor.multisample.count = 1;

    WGPUBindGroupLayoutDescriptor bind_group_layout_descriptor = {};
    bind_group_layout_descriptor.entryCount = group_entries.size();
    bind_group_layout_descriptor.entries = group_entries.data();

    auto layout = wgpuDeviceCreateBindGroupLayout(device, &bind_group_layout_descriptor);

    WGPUPipelineLayoutDescriptor pipeline_layout_descriptor = {};
    pipeline_layout_descriptor.label = pipeline->label.c_str();
    pipeline_layout_descriptor.bindGroupLayoutCount = 1;
    pipeline_layout_descriptor.bindGroupLayouts = &layout;

    pipeline->layout =  wgpuDeviceCreatePipelineLayout(device, &pipeline_layout_descriptor);

    pipeline->render_handle = wgpuDeviceCreateRenderPipeline(device, &descriptor);
    pipeline->bind_group_layout = wgpuRenderPipelineGetBindGroupLayout(pipeline->render_handle, 0);

    return pipeline;
}

GFXPipeline* GFXWebGPU::create_compute_pipeline(const GFXComputePipelineCreateInfo& info) {
    auto pipeline = new GFXWebGPUPipeline();

    WGPUComputePipelineDescriptor descriptor = {};
    descriptor.label = info.label.c_str();

    {
        const bool use_shader_source = !info.compute_src.is_path();

        if (use_shader_source) {
            auto compute_shader_vector = info.compute_src.as_bytecode();

            descriptor.compute.module = create_shader(compute_shader_vector.data(), compute_shader_vector.size() * sizeof(uint32_t), info.label.c_str());
        }
        else {
            auto compute_shader = prism::open_file(prism::internal_domain / (info.compute_src.as_path().string() + ".wgsl"), true);
            compute_shader->read_all();

            descriptor.compute.module = create_shader(compute_shader->cast_data<uint32_t>(), compute_shader->size(), info.compute_src.as_path().string().c_str());
        }
    }


    pipeline->compute_handle = wgpuDeviceCreateComputePipeline(device, &descriptor);

    WGPUBindGroupLayoutDescriptor bind_group_layout_descriptor = {};

    pipeline->bind_group_layout = wgpuDeviceCreateBindGroupLayout(device, &bind_group_layout_descriptor);

    return pipeline;
}

void GFXWebGPU::submit(GFXCommandBuffer* command_buffer, const platform::window_ptr window) {
    WGPUTextureView backBufView = nullptr;
    if(window != nullptr)
        backBufView = wgpuSwapChainGetCurrentTextureView(swapchain);

    WGPUCommandEncoder command_encoder = wgpuDeviceCreateCommandEncoder(device, nullptr);

    GFXWebGPUPipeline* current_pipeline = nullptr;
    WGPURenderPassEncoder render_encoder = nullptr;
    WGPUComputePassEncoder compute_encoder = nullptr;
    GFXWebGPURenderPass* current_render_pass = nullptr;
    GFXWebGPUFramebuffer* current_framebuffer = nullptr;
    WGPUColor current_clear_color = {};
    Viewport current_viewport {}; // lol webgpu doesn't even have a viewport type??

    enum class CurrentEncoder {
        None,
        Render,
        Compute,
    } current_encoder = CurrentEncoder::None;

    const auto need_encoder = [&](CurrentEncoder encoder, bool needs_reset = false) {
        if(encoder != current_encoder || needs_reset) {
            if(render_encoder != nullptr)
                wgpuRenderPassEncoderEndPass(render_encoder);

            if(compute_encoder != nullptr)
                wgpuComputePassEncoderEndPass(compute_encoder);

            render_encoder = nullptr;
            compute_encoder = nullptr;
        }

        if(current_encoder == encoder && !needs_reset)
            return;

        switch(encoder) {
            case CurrentEncoder::None:
                break;
            case CurrentEncoder::Render:
            {
                WGPURenderPassDescriptor render_pass_descriptor = {};

                std::vector<WGPURenderPassColorAttachment> color_attachments;

                if(current_framebuffer != nullptr) {
                    unsigned int i = 0;
                    for(const auto& attachment : current_framebuffer->attachments) {
                        if(attachment->format == WGPUTextureFormat_Depth32Float) {
                            // TODO: lol what
                            auto depth_attachment = new WGPURenderPassDepthStencilAttachment();
                            depth_attachment->view = attachment->view;
                            depth_attachment->depthLoadOp = WGPULoadOp_Clear;
                            depth_attachment->depthStoreOp = WGPUStoreOp_Store;

                            render_pass_descriptor.depthStencilAttachment = depth_attachment;
                        } else {
                            WGPURenderPassColorAttachment color_attachment = {};
                            color_attachment.view = attachment->view;
                            color_attachment.loadOp = WGPULoadOp_Clear;
                            color_attachment.storeOp = WGPUStoreOp_Store;
                            color_attachment.clearColor = current_clear_color;

                            color_attachments.push_back(color_attachment);
                        }
                    }
                } else {
                    // we are rendering to the screen
                    WGPURenderPassColorAttachment color_attachment = {};
                    color_attachment.view = backBufView;
                    color_attachment.loadOp = WGPULoadOp_Clear;
                    color_attachment.storeOp = WGPUStoreOp_Store;
                    color_attachment.clearColor = current_clear_color;

                    color_attachments.push_back(color_attachment);
                }

                render_pass_descriptor.colorAttachmentCount = color_attachments.size();
                render_pass_descriptor.colorAttachments = color_attachments.data();

                render_encoder = wgpuCommandEncoderBeginRenderPass(command_encoder, &render_pass_descriptor);

                //if(currentViewport.width != 0.0f && currentViewport.height != 0.0f)
                //    renderEncoder->setViewport(currentViewport);
            }
                break;
            case CurrentEncoder::Compute:
            {
                WGPUComputePassDescriptor compute_pass_descriptor = {};

                compute_encoder = wgpuCommandEncoderBeginComputePass(command_encoder, &compute_pass_descriptor);
            }
                break;
        }

        current_encoder = encoder;
    };

    uint64_t last_bind_group_hash = 0;

    const auto try_bind_group = [&]() -> bool {
        if(current_pipeline == nullptr)
            return false;

        const uint64_t hash = get_bind_group_hash(current_pipeline);
        if(last_bind_group_hash != hash) {
            if(!current_pipeline->cached_bind_groups.count(hash))
                cache_bind_group_state(current_pipeline, current_pipeline->bind_group_layout);

            auto bind_group = current_pipeline->cached_bind_groups[hash];
            if(bind_group == nullptr)
                return false;

            wgpuRenderPassEncoderSetBindGroup(render_encoder, 0, bind_group, 0, nullptr);

            last_bind_group_hash = hash;
        }

        return true;
    };

    for(auto command : command_buffer->commands) {
        switch (command.type) {
            case GFXCommandType::Invalid:
                break;
            case GFXCommandType::SetRenderPass:
            {
                current_clear_color = {};
                current_clear_color.r = command.data.set_render_pass.clear_color.r;
                current_clear_color.g = command.data.set_render_pass.clear_color.g;
                current_clear_color.b = command.data.set_render_pass.clear_color.b;
                current_clear_color.a = command.data.set_render_pass.clear_color.a;

                current_framebuffer = (GFXWebGPUFramebuffer*)command.data.set_render_pass.framebuffer;
                current_render_pass = (GFXWebGPURenderPass*)command.data.set_render_pass.render_pass;
                current_viewport = {};

                need_encoder(CurrentEncoder::Render, true);
            }
                break;
            case GFXCommandType::EndRenderPass:
            {
                current_render_pass = nullptr;

                if(current_encoder == CurrentEncoder::Render)
                    wgpuRenderPassEncoderEndPass(render_encoder);

                render_encoder = nullptr;
            }
            case GFXCommandType::SetGraphicsPipeline:
            {
                need_encoder(CurrentEncoder::Render);

                current_pipeline = (GFXWebGPUPipeline*)command.data.set_graphics_pipeline.pipeline;
                if(current_pipeline != nullptr) {
                    wgpuRenderPassEncoderSetPipeline(render_encoder, current_pipeline->render_handle);

                    reset_bind_state();
                    last_bind_group_hash = 0;
                }
            }
                break;
            case GFXCommandType::SetComputePipeline:
            {
                need_encoder(CurrentEncoder::Compute);

                current_pipeline = (GFXWebGPUPipeline*)command.data.set_compute_pipeline.pipeline;
                if(current_pipeline != nullptr) {
                    wgpuComputePassEncoderSetPipeline(compute_encoder, current_pipeline->compute_handle);

                    reset_bind_state();
                    last_bind_group_hash = 0;
                }
            }
                break;
            case GFXCommandType::SetVertexBuffer:
            {
                need_encoder(CurrentEncoder::Render);

                wgpuRenderPassEncoderSetVertexBuffer(render_encoder,
                                                     command.data.set_vertex_buffer.index,
                                                     ((GFXWebGPUBuffer*)command.data.set_vertex_buffer.buffer)->handle,
                                                     command.data.set_vertex_buffer.offset,
                                                     ((GFXWebGPUBuffer*)command.data.set_vertex_buffer.buffer)->size);
            }
                break;
            case GFXCommandType::SetIndexBuffer:
            {
                need_encoder(CurrentEncoder::Render);

                wgpuRenderPassEncoderSetIndexBuffer(render_encoder,
                                                    ((GFXWebGPUBuffer*)command.data.set_index_buffer.buffer)->handle,
                                                    command.data.set_index_buffer.index_type == IndexType::UINT32 ? WGPUIndexFormat_Uint32 : WGPUIndexFormat_Uint16,
                                                    0,
                                                    ((GFXWebGPUBuffer*)command.data.set_index_buffer.buffer)->size);
            }
                break;
            case GFXCommandType::BindShaderBuffer:
            {
                BoundShaderBuffer bsb;
                bsb.buffer = command.data.bind_shader_buffer.buffer;
                bsb.offset = command.data.bind_shader_buffer.offset;
                bsb.size = command.data.bind_shader_buffer.size;

                boundShaderBuffers[command.data.bind_shader_buffer.index] = bsb;
            }
                break;
            case GFXCommandType::BindTexture:
            {
                boundTextures[command.data.bind_texture.index] = command.data.bind_texture.texture;
            }
                break;
            case GFXCommandType::BindSampler:
            {
                boundSamplers[command.data.bind_sampler.index] = command.data.bind_sampler.sampler;
            }
                break;
            case GFXCommandType::SetPushConstant:
            {
                need_encoder(CurrentEncoder::Render);

                if(current_pipeline != nullptr) {
                    uint32_t offset = current_pipeline->current_push_constant_offset;

                    wgpuQueueWriteBuffer(queue, current_pipeline->push_constant_emulation_buffer->handle, offset, command.data.set_push_constant.bytes.data(), command.data.set_push_constant.bytes.size());

                    current_pipeline->current_push_constant_offset += 128;

                    BoundShaderBuffer bsb;
                    bsb.buffer = current_pipeline->push_constant_emulation_buffer;
                    bsb.offset = offset;
                    bsb.size = command.data.set_push_constant.size;

                    boundShaderBuffers[current_pipeline->push_constant_emulation_buffer_index] = bsb;
                }
            }
                break;
            case GFXCommandType::Draw:
            {
                need_encoder(CurrentEncoder::Render);

                if(current_pipeline == nullptr)
                    continue;

                if(try_bind_group()) {
                    wgpuRenderPassEncoderDraw(render_encoder,
                                              command.data.draw.vertex_count,
                                              command.data.draw.instance_count,
                                              command.data.draw.vertex_offset,
                                              command.data.draw.base_instance);
                }
            }
                break;
            case GFXCommandType::DrawIndexed:
            {
                need_encoder(CurrentEncoder::Render);

                if(current_pipeline == nullptr)
                    continue;

                if(try_bind_group()) {
                    wgpuRenderPassEncoderDrawIndexed(render_encoder,
                                                 command.data.draw_indexed.index_count,
                                                 1,
                                                 command.data.draw_indexed.vertex_offset,
                                                 command.data.draw_indexed.vertex_offset,
                                                 command.data.draw_indexed.base_instance);
                }
            }
                break;
            case GFXCommandType::MemoryBarrier:
                // not supported
                break;
            case GFXCommandType::CopyTexture:
            {
                // TODO: blit op
            }
                break;
            case GFXCommandType::SetViewport:
            {
                need_encoder(CurrentEncoder::Render);

                current_viewport = command.data.set_viewport.viewport;

                wgpuRenderPassEncoderSetViewport(render_encoder,
                                                 current_viewport.x,
                                                 current_viewport.y,
                                                 current_viewport.width,
                                                 current_viewport.height,
                                                 current_viewport.min_depth,
                                                 current_viewport.max_depth);
            }
                break;
            case GFXCommandType::SetScissor:
            {
                need_encoder(CurrentEncoder::Render);

                wgpuRenderPassEncoderSetScissorRect(render_encoder,
                                                    command.data.set_scissor.rect.offset.x,
                                                    command.data.set_scissor.rect.offset.y,
                                                    command.data.set_scissor.rect.extent.width,
                                                    command.data.set_scissor.rect.extent.height);
            }
                break;
            case GFXCommandType::GenerateMipmaps:
            {
                // TODO: not supported by webgpu?
            }
                break;
            case GFXCommandType::SetDepthBias: {
                need_encoder(CurrentEncoder::Render);

                // TODO: not supported by webgpu?
            }
                break;
            case GFXCommandType::PushGroup:
            {
                // TOOD: stub
            }
                break;
            case GFXCommandType::PopGroup:
            {
                // TOOD: stub
            }
                break;
            case GFXCommandType::InsertLabel:
            {
                // TOOD: stub
            }
                break;
            case GFXCommandType::Dispatch:
            {
                need_encoder(CurrentEncoder::Compute);

                if(try_bind_group()) {
                    wgpuComputePassEncoderDispatch(compute_encoder,
                                                   command.data.dispatch.group_count_x,
                                                   command.data.dispatch.group_count_y,
                                                   command.data.dispatch.group_count_z);
                }
            }
                break;
            default:
                prism::log("Unhandled GFX command {}", utility::enum_to_string(command.type));
        }
    }

    if(render_encoder != nullptr)
        wgpuRenderPassEncoderEndPass(render_encoder);

    if(compute_encoder != nullptr)
        wgpuComputePassEncoderEndPass(compute_encoder);

    WGPUCommandBuffer commands = wgpuCommandEncoderFinish(command_encoder, nullptr);

    wgpuQueueSubmit(queue, 1, &commands);
    wgpuCommandBufferRelease(commands);

    if(window != nullptr) {
        wgpuTextureViewRelease(backBufView);
    }
}

uint64_t GFXWebGPU::get_bind_group_hash(GFXWebGPUPipeline *pipeline) {
    uint64_t hash = 0;
    hash += (int64_t)pipeline;

    int i = 0;
    for (auto& buffer : boundShaderBuffers) {
        if (buffer.buffer != nullptr) {
            hash += (uint64_t)buffer.buffer * (i + 1);
        }
    }

    i = 0;
    for (auto& texture : boundTextures) {
        if (texture != nullptr) {
            hash += (uint64_t)texture * (i + 1);
        }
    }

    return hash;
}

void GFXWebGPU::cache_bind_group_state(GFXWebGPUPipeline* pipeline, WGPUBindGroupLayout group_layout) {
    const uint64_t hash = get_bind_group_hash(pipeline);

    std::vector<WGPUBindGroupEntry> group_entries;

    prism::log("creating bind group for {}", pipeline->label);

    for (auto [i, buffer] : utility::enumerate(boundShaderBuffers)) {
        if (buffer.buffer != nullptr) {
            auto wgpu_buffer = (GFXWebGPUBuffer*)buffer.buffer;

            prism::log("binding {} filled", i);

            WGPUBindGroupEntry entry = {};
            entry.buffer = wgpu_buffer->handle;
            entry.size = buffer.size;
            entry.offset = buffer.offset;
            entry.binding = i;

            group_entries.push_back(entry);
        }
    }

    for (auto [i, texture] : utility::enumerate(boundTextures)) {
        if (texture != nullptr) {
            auto wgpu_texture = (GFXWebGPUTexture*) texture;

            prism::log("binding {} filled", i);

            WGPUBindGroupEntry entry = {};
            entry.textureView = wgpu_texture->view;
            entry.binding = i;

            group_entries.push_back(entry);
        }
    }

    for (auto [i, sampler] : utility::enumerate(boundSamplers)) {
        if (sampler != nullptr) {
            auto wgpu_sampler = (GFXWebGPUSampler*) sampler;

            prism::log("binding {} filled", i);

            WGPUBindGroupEntry entry = {};
            entry.sampler = wgpu_sampler->handle;
            entry.binding = i;

            group_entries.push_back(entry);
        }
    }

    prism::log("num entries: {}", group_entries.size());

    prism::log("-----");

    WGPUBindGroupDescriptor group_descriptor = {};
    group_descriptor.layout = group_layout;
    group_descriptor.entryCount = group_entries.size();
    group_descriptor.entries = group_entries.data();
    group_descriptor.label = pipeline->label.c_str();

    pipeline->cached_bind_groups[hash] = wgpuDeviceCreateBindGroup(device, &group_descriptor);
}

void GFXWebGPU::reset_bind_state() {
    for (auto& buffer : boundShaderBuffers)
        buffer.buffer = nullptr;

    for (auto& texture : boundTextures)
        texture = nullptr;

    for (auto& sampler : boundSamplers)
        sampler = nullptr;
}

bool GFXWebGPU::supports_feature(GFXFeature feature) {
    return true;
}