const std = @import("std");
const au = @import("au.zig");
const vk = @import("vk");
const shaders = @import("shaders");

const Self = @This();

set_layout: vk.DescriptorSetLayout,
layout: vk.PipelineLayout,
pipeline: vk.Pipeline,

pub const Index = u16;

pub const Uniform = extern struct {
    mat: [16]f32,

    const Bindings = [_]vk.DescriptorSetLayoutBinding{.{
        .binding = 0,
        .descriptor_type = .uniform_buffer,
        .descriptor_count = 1,
        .stage_flags = .{ .vertex_bit = true },
    }};
};

pub const Vertex = extern struct {
    pos: [4]f32,
    color: [3]f32,

    const Bindings = [_]vk.VertexInputBindingDescription{
        .{
            .binding = 0,
            .stride = @sizeOf(Vertex),
            .input_rate = .vertex,
        },
    };

    const Attributes = [_]vk.VertexInputAttributeDescription{
        .{
            .binding = 0,
            .location = 0,
            .format = .r32g32b32a32_sfloat,
            .offset = @offsetOf(Vertex, "pos"),
        },
        .{
            .binding = 0,
            .location = 1,
            .format = .r32g32b32_sfloat,
            .offset = @offsetOf(Vertex, "color"),
        },
    };
};

pub fn init(cache: vk.PipelineCache) !Self {
    const vert = try au.D.createShaderModule(&.{
        .code_size = shaders.triangle_vert.len,
        .p_code = @ptrCast(&shaders.triangle_vert),
    }, null);
    defer au.D.destroyShaderModule(vert, null);

    const frag = try au.D.createShaderModule(&.{
        .code_size = shaders.triangle_frag.len,
        .p_code = @ptrCast(&shaders.triangle_frag),
    }, null);
    defer au.D.destroyShaderModule(frag, null);

    const set_layout = try au.D.createDescriptorSetLayout(&vk.DescriptorSetLayoutCreateInfo{
        .flags = .{},
        .binding_count = @intCast(Uniform.Bindings.len),
        .p_bindings = &Uniform.Bindings,
    }, null);
    errdefer au.D.destroyDescriptorSetLayout(set_layout, null);

    const layout = try au.D.createPipelineLayout(&vk.PipelineLayoutCreateInfo{
        .push_constant_range_count = 0,
        .set_layout_count = 1,
        .p_set_layouts = &.{set_layout},
    }, null);
    errdefer au.D.destroyPipelineLayout(layout, null);

    var pipeline: vk.Pipeline = .null_handle;
    _ = try au.D.createGraphicsPipelines(cache, 1, &[1]vk.GraphicsPipelineCreateInfo{
        vk.GraphicsPipelineCreateInfo{
            .stage_count = 2,
            .p_stages = &.{
                vk.PipelineShaderStageCreateInfo{ .stage = .{ .vertex_bit = true }, .module = vert, .p_name = "main" },
                vk.PipelineShaderStageCreateInfo{ .stage = .{ .fragment_bit = true }, .module = frag, .p_name = "main" },
            },
            .layout = layout,
            .render_pass = .null_handle,
            .subpass = 0,
            .base_pipeline_handle = .null_handle,
            .base_pipeline_index = -1,
            .p_vertex_input_state = &vk.PipelineVertexInputStateCreateInfo{
                .vertex_binding_description_count = @intCast(Vertex.Bindings.len),
                .p_vertex_binding_descriptions = &Vertex.Bindings,
                .vertex_attribute_description_count = @intCast(Vertex.Attributes.len),
                .p_vertex_attribute_descriptions = &Vertex.Attributes,
            },
            .p_input_assembly_state = &vk.PipelineInputAssemblyStateCreateInfo{
                .topology = .triangle_list,
                .primitive_restart_enable = vk.FALSE,
            },
            .p_tessellation_state = null,
            .p_viewport_state = &vk.PipelineViewportStateCreateInfo{
                .viewport_count = 1,
                .scissor_count = 1,
            },
            .p_rasterization_state = &vk.PipelineRasterizationStateCreateInfo{
                .depth_clamp_enable = vk.FALSE,
                .rasterizer_discard_enable = vk.FALSE,
                .polygon_mode = .fill,
                .cull_mode = .{ .back_bit = true },
                .front_face = .counter_clockwise,
                .depth_bias_enable = vk.FALSE,
                .depth_bias_constant_factor = 0.0,
                .depth_bias_clamp = 0.0,
                .depth_bias_slope_factor = 0.0,
                .line_width = 1.0,
            },
            .p_multisample_state = &vk.PipelineMultisampleStateCreateInfo{
                .rasterization_samples = .{ .@"1_bit" = true },
                .sample_shading_enable = vk.FALSE,
                .min_sample_shading = 1,
                .alpha_to_coverage_enable = vk.FALSE,
                .alpha_to_one_enable = vk.FALSE,
            },
            .p_depth_stencil_state = null,
            .p_color_blend_state = &vk.PipelineColorBlendStateCreateInfo{
                .logic_op_enable = vk.FALSE,
                .logic_op = .copy,
                .blend_constants = [_]f32{ 0, 0, 0, 0 },
                .attachment_count = 1,
                .p_attachments = &.{
                    vk.PipelineColorBlendAttachmentState{
                        .blend_enable = vk.FALSE,
                        .color_blend_op = .add,
                        .src_color_blend_factor = .one,
                        .dst_color_blend_factor = .zero,
                        .alpha_blend_op = .add,
                        .src_alpha_blend_factor = .one,
                        .dst_alpha_blend_factor = .zero,
                        .color_write_mask = .{ .r_bit = true, .g_bit = true, .b_bit = true, .a_bit = true },
                    },
                },
            },
            .p_dynamic_state = &vk.PipelineDynamicStateCreateInfo{
                .flags = .{},
                .dynamic_state_count = 2,
                .p_dynamic_states = &.{
                    .viewport,
                    .scissor,
                },
            },
            .p_next = &vk.PipelineRenderingCreateInfo{
                .color_attachment_count = 1,
                .p_color_attachment_formats = &.{au.device_config.format.format},
                .depth_attachment_format = .undefined,
                .stencil_attachment_format = .undefined,
                .view_mask = 0,
            },
        },
    }, null, @ptrCast(&pipeline));
    errdefer au.D.destroyPipeline(pipeline, null);

    return .{ .pipeline = pipeline, .layout = layout, .set_layout = set_layout };
}

pub fn deinit(self: Self) void {
    au.D.destroyPipeline(self.pipeline, null);
    au.D.destroyPipelineLayout(self.layout, null);
    au.D.destroyDescriptorSetLayout(self.set_layout, null);
}