//! ----------------------------------------------------
//! ----------------------------------------------------

const std = @import("std");
const gtl = @import("gtl");
const wgpu = @import("wgpu");
const Self = @This();

const Material = @import("../../common/component/material.zig");
const ShaderCache = @import("shader_cache.zig");
const ShaderID = ShaderCache.ShaderID;
const Vertex = @import("mesh_cache.zig").Vertex;

/// ----------------------------------------------------
/// ----------------------------------------------------
pub const GPUPipeline = struct {
    raw: *wgpu.WGPURenderPipelineImpl,
    layout: *wgpu.WGPUPipelineLayoutImpl,
};

/// ----------------------------------------------------
/// ----------------------------------------------------
pub const GPUMaterial = struct {
    buffer: *wgpu.WGPUBufferImpl,
    bind_group: *wgpu.WGPUBindGroupImpl,
};

//
// FIELDS
//

pipelines: gtl.KeyedSlotMap(Material.Unique, GPUPipeline),
materials: gtl.SlotMap(Material),
materials_gpu: gtl.KeyedSlotMap(gtl.Handle(Material), GPUMaterial),
shader_cache: *ShaderCache,
queue: *wgpu.WGPUQueueImpl,
device: *wgpu.WGPUDeviceImpl,
alloc: std.mem.Allocator,

/// ----------------------------------------------------
/// ----------------------------------------------------
pub fn init(
    shader_cache: *ShaderCache,
    queue: *wgpu.WGPUQueueImpl,
    device: *wgpu.WGPUDeviceImpl,
    alloc: std.mem.Allocator,
) Self {
    return .{
        .pipelines = .init(alloc),
        .materials = .init(alloc),
        .materials_gpu = .init(alloc),
        .shader_cache = shader_cache,
        .queue = queue,
        .device = device,
        .alloc = alloc,
    };
}

/// ----------------------------------------------------
/// ----------------------------------------------------
pub fn deinit(self: *Self) void {
    // --- FREE PIPELINES ---
    for (self.pipelines.items()) |entry| {
        // wgpu.wgpuPipelineLayoutRelease(entry.value.layout);
        wgpu.wgpuRenderPipelineRelease(entry.value.raw);
    }

    // --- FREE MATERIALS ---
    for (self.materials_gpu.items()) |entry| {
        wgpu.wgpuBufferRelease(entry.value.buffer);
        wgpu.wgpuBindGroupRelease(entry.value.bind_group);
    }

    // --- # ---
    self.pipelines.deinit();
    self.materials.deinit();
    self.materials_gpu.deinit();
}

/// ----------------------------------------------------
/// ----------------------------------------------------
pub fn getOrCreateMaterial(
    self: *Self,
    material: Material,
) !*GPUMaterial {
    return try self.materials.getOrPut(
        material.unique.shader,
        blk: {
            // --- UNIFORM ---
            const uniform = wgpu.wgpuDeviceCreateBuffer(self.device, &wgpu.WGPUBufferDescriptor{
                .usage = wgpu.WGPUBufferUsage_CopyDst | wgpu.WGPUBufferUsage_Uniform,
                .size = @sizeOf(Material.Style),
            }) orelse return error.Buffer;

            wgpu.wgpuQueueWriteBuffer(
                self.queue,
                uniform,
                0,
                &material.style,
                @sizeOf(Material.Style),
            );

            // --- RESULT ---
            break :blk .{
                .bind_group = undefined,
                .buffer = uniform,
            };
        },
    );
}

/// ----------------------------------------------------
/// ----------------------------------------------------
pub fn getOrCreate(
    self: *Self,
    key: Material,
) !*GPUPipeline {
    return try self.pipelines.getOrPut(key.unique, blk: {
        // const layout = wgpu.wgpuDeviceCreatePipelineLayout(
        //     self.device,
        //     &wgpu.WGPUPipelineLayoutDescriptor{
        //         .bindGroupLayoutCount = 1,
        //         .bindGroupLayouts = &wgpu.WGPUBindGroupLayoutDescriptor{
        //             .entryCount = 1,
        //         },
        //     },
        // ) orelse return error.Layout;

        // --- RESULT ---
        break :blk .{
            .raw = try self.createPipeline(key.unique, undefined),
            .layout = undefined,
        };
    });
}

/// ----------------------------------------------------
/// ----------------------------------------------------
fn createPipeline(
    self: *Self,
    key: Material.Unique,
    layout: *wgpu.WGPUPipelineLayoutImpl,
) !*wgpu.WGPURenderPipelineImpl {
    _ = layout;

    // --- ENSURE SHADER MODULE ---
    const shader_module = try self.shader_cache.getOrCreate(
        key.shader,
    );

    // --- RESULT ---
    return wgpu.wgpuDeviceCreateRenderPipeline(self.device, &wgpu.WGPURenderPipelineDescriptor{
        .label = .{ .data = @tagName(key.shader), .length = wgpu.WGPU_STRLEN },
        .vertex = .{
            .module = shader_module.raw,
            .entryPoint = .{ .data = "vs_main", .length = wgpu.WGPU_STRLEN },
            .bufferCount = 1,
            .buffers = &wgpu.WGPUVertexBufferLayout{
                .attributeCount = 3,
                .attributes = &[_]wgpu.WGPUVertexAttribute{ // TODO:  Mabe mabe allow for custom, mabe mabe later tho
                    .{
                        .format = wgpu.WGPUVertexFormat_Float32x3,
                        .offset = @offsetOf(Vertex, "position"),
                        .shaderLocation = 0,
                    },
                    .{
                        .format = wgpu.WGPUVertexFormat_Float32x3,
                        .offset = @offsetOf(Vertex, "normal"),
                        .shaderLocation = 1,
                    },
                    .{
                        .format = wgpu.WGPUVertexFormat_Float32x2,
                        .offset = @offsetOf(Vertex, "uv"),
                        .shaderLocation = 2,
                    },
                },
                .arrayStride = @sizeOf(Vertex),
                .stepMode = wgpu.WGPUVertexStepMode_Vertex,
            },
        },
        .fragment = &wgpu.WGPUFragmentState{
            .module = shader_module.raw,
            .entryPoint = .{ .data = "fs_main", .length = wgpu.WGPU_STRLEN },
            .targetCount = 1,
            .targets = &wgpu.WGPUColorTargetState{
                .format = 28, // FIX: Use proper format
                .blend = &wgpu.WGPUBlendState{
                    .color = .{
                        .srcFactor = wgpu.WGPUBlendFactor_SrcAlpha,
                        .dstFactor = wgpu.WGPUBlendFactor_OneMinusSrcAlpha,
                        .operation = wgpu.WGPUBlendOperation_Add,
                    },
                    .alpha = .{
                        .srcFactor = wgpu.WGPUBlendFactor_Zero,
                        .dstFactor = wgpu.WGPUBlendFactor_One,
                        .operation = wgpu.WGPUBlendOperation_Add,
                    },
                },
                .writeMask = wgpu.WGPUColorWriteMask_All,
            },
        },
        .primitive = .{
            .topology = switch (key.topology) {
                .triangle_list => wgpu.WGPUPrimitiveTopology_TriangleList,
                .point_list => wgpu.WGPUPrimitiveTopology_PointList,
                .line_list => wgpu.WGPUPrimitiveTopology_LineList,
            },
            .frontFace = switch (key.front_face) {
                .ccw => wgpu.WGPUFrontFace_CCW,
                .cw => wgpu.WGPUFrontFace_CW,
            },
            .cullMode = switch (key.cull_mode) {
                .back => wgpu.WGPUCullMode_Back,
                .front => wgpu.WGPUCullMode_Front,
                .none => wgpu.WGPUCullMode_None,
            },
        },
        .multisample = .{
            .count = 1,
            .mask = ~@as(u32, 0),
            .alphaToCoverageEnabled = 0,
        },
    }) orelse {
        return error.PipelineCreationFailed;
    };
}