#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* GFXWebGPU::acquire_command_buffer(bool for_presentation_use) {
    return new GFXCommandBuffer();
}

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

    buffer->size = size;

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

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

    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;

    wgpuBufferUnmap(wgpu_buffer->handle);
}

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

    WGPUTextureDescriptor descriptor = {};
    descriptor.size.width = info.width;
    descriptor.size.height = info.height;
    descriptor.size.depthOrArrayLayers = info.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;

    texture->handle = wgpuDeviceCreateTexture(device, &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();

    return sampler;
}

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

    return framebuffer;
}

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

    return render_pass;
}

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

    WGPURenderPipelineDescriptor descriptor = {};
    descriptor.label = info.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()), 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());
        }
    }

    WGPUFragmentState fragment = {};

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

        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()), 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());
        }
    }

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

    // 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;

        prism::log("{} of {}", attribute.binding, info.vertex_input.inputs.size());

        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);

        prism::log("debugging attributes at loc {}", binding.location);

        for(auto attr : attributes[binding.location]) {
            prism::log("- attrib {}", attr.shaderLocation);
            prism::log("fmt: {}", utility::enum_to_string(attr.format));
        }

        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.stripIndexFormat = WGPUIndexFormat_Uint16;
    descriptor.primitive.topology = WGPUPrimitiveTopology_TriangleList;

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

    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()), 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);

    return pipeline;
}

void GFXWebGPU::submit(GFXCommandBuffer* command_buffer, const platform::window_ptr window) {
    WGPUTextureView backBufView = wgpuSwapChainGetCurrentTextureView(swapchain);            // create textureView

    WGPUCommandEncoder encoder = wgpuDeviceCreateCommandEncoder(device, nullptr);
    WGPUCommandBuffer commands = wgpuCommandEncoderFinish(encoder, nullptr);

    wgpuQueueSubmit(queue, 1, &commands);
    wgpuCommandBufferRelease(commands);
    wgpuTextureViewRelease(backBufView);
}