package pipeline

import bk "../backend"
import gpu "../core"
import "core:log"
import vk "vendor:vulkan"

Graphics_Pipeline :: struct {
	pipeline: vk.Pipeline,
	layout:   vk.PipelineLayout,
}

Pipeline_Config :: struct {
	vert_module:            vk.ShaderModule,
	frag_module:            vk.ShaderModule,
	render_pass:            vk.RenderPass,
	topology:               vk.PrimitiveTopology,
	polygon_mode:           vk.PolygonMode,
	cull_mode:              vk.CullModeFlags,
	front_face:             vk.FrontFace,
	enable_blending:        bool,
	blend_mode:             bk.Blend_Mode,
	enable_depth_test:      bool,
	enable_depth_bias:      bool,
	stencil:                bk.Stencil_State,
	depth_only:             bool, // no color attachment (shadow pass)
	color_attachment_count: u32,
	push_constant_size:     u32,
	push_constant_stages:   vk.ShaderStageFlags,
	descriptor_set_layouts: []vk.DescriptorSetLayout,
	binding_descriptions:   []vk.VertexInputBindingDescription,
	attribute_descriptions: []vk.VertexInputAttributeDescription,
}

create_graphics_pipeline :: proc(
	dev: ^gpu.Gpu_Device,
	config: Pipeline_Config,
) -> (
	gp: Graphics_Pipeline,
	ok: bool,
) {
	// Shader stages
	stages := [2]vk.PipelineShaderStageCreateInfo {
		{
			sType = .PIPELINE_SHADER_STAGE_CREATE_INFO,
			stage = {.VERTEX},
			module = config.vert_module,
			pName = "main",
		},
		{
			sType = .PIPELINE_SHADER_STAGE_CREATE_INFO,
			stage = {.FRAGMENT},
			module = config.frag_module,
			pName = "main",
		},
	}

	// Vertex input
	vertex_input := vk.PipelineVertexInputStateCreateInfo {
		sType                           = .PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO,
		vertexBindingDescriptionCount   = u32(len(config.binding_descriptions)),
		pVertexBindingDescriptions      = raw_data(config.binding_descriptions),
		vertexAttributeDescriptionCount = u32(len(config.attribute_descriptions)),
		pVertexAttributeDescriptions    = raw_data(config.attribute_descriptions),
	}

	// Input assembly
	input_assembly := vk.PipelineInputAssemblyStateCreateInfo {
		sType                  = .PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO,
		topology               = config.topology,
		primitiveRestartEnable = false,
	}

	// Dynamic state: viewport + scissor (+ depth bias if enabled)
	dynamic_states: [3]vk.DynamicState
	dynamic_states[0] = .VIEWPORT
	dynamic_states[1] = .SCISSOR
	dynamic_state_count: u32 = 2
	if config.enable_depth_bias {
		dynamic_states[2] = .DEPTH_BIAS
		dynamic_state_count = 3
	}
	dynamic_state := vk.PipelineDynamicStateCreateInfo {
		sType             = .PIPELINE_DYNAMIC_STATE_CREATE_INFO,
		dynamicStateCount = dynamic_state_count,
		pDynamicStates    = &dynamic_states[0],
	}

	// Viewport state (counts only, actual values set dynamically)
	viewport_state := vk.PipelineViewportStateCreateInfo {
		sType         = .PIPELINE_VIEWPORT_STATE_CREATE_INFO,
		viewportCount = 1,
		scissorCount  = 1,
	}

	// Rasterization
	rasterization := vk.PipelineRasterizationStateCreateInfo {
		sType                   = .PIPELINE_RASTERIZATION_STATE_CREATE_INFO,
		depthClampEnable        = false,
		rasterizerDiscardEnable = false,
		polygonMode             = config.polygon_mode,
		cullMode                = config.cull_mode,
		frontFace               = config.front_face,
		depthBiasEnable         = b32(config.enable_depth_bias),
		lineWidth               = 1.0,
	}

	// Multisample
	multisample := vk.PipelineMultisampleStateCreateInfo {
		sType                = .PIPELINE_MULTISAMPLE_STATE_CREATE_INFO,
		rasterizationSamples = {._1},
		sampleShadingEnable  = false,
	}

	// Depth stencil
	depth_stencil := vk.PipelineDepthStencilStateCreateInfo {
		sType                 = .PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO,
		depthTestEnable       = b32(config.enable_depth_test),
		depthWriteEnable      = b32(config.enable_depth_test),
		depthCompareOp        = .LESS,
		depthBoundsTestEnable = false,
		stencilTestEnable     = b32(config.stencil.enable),
		front                 = to_vk_stencil_face(config.stencil.front, config.stencil),
		back                  = to_vk_stencil_face(config.stencil.back, config.stencil),
	}

	// Color blend attachments
	blend_attachments: [bk.MAX_COLOR_TARGETS]vk.PipelineColorBlendAttachmentState
	attachment_count := config.color_attachment_count
	if !config.depth_only && attachment_count == 0 {
		attachment_count = 1
	}
	for i in 0..<attachment_count {
		blend_attachments[i] = vk.PipelineColorBlendAttachmentState {
			colorWriteMask = {.R, .G, .B, .A},
		}
	}
	if config.enable_blending {
		factors := bk.blend_factors(config.blend_mode)
		for i in 0..<attachment_count {
			blend_attachments[i].blendEnable = true
			blend_attachments[i].colorBlendOp = .ADD
			blend_attachments[i].srcColorBlendFactor = to_vk_blend_factor(factors.src_color)
			blend_attachments[i].dstColorBlendFactor = to_vk_blend_factor(factors.dst_color)
			blend_attachments[i].srcAlphaBlendFactor = to_vk_blend_factor(factors.src_alpha)
			blend_attachments[i].dstAlphaBlendFactor = to_vk_blend_factor(factors.dst_alpha)
			blend_attachments[i].alphaBlendOp = .ADD
		}
	}

	color_blend := vk.PipelineColorBlendStateCreateInfo {
		sType           = .PIPELINE_COLOR_BLEND_STATE_CREATE_INFO,
		logicOpEnable   = false,
		attachmentCount = config.depth_only ? 0 : attachment_count,
		pAttachments    = config.depth_only ? nil : &blend_attachments[0],
	}

	// Pipeline layout with optional push constants
	push_range := vk.PushConstantRange {
		stageFlags = config.push_constant_stages,
		offset     = 0,
		size       = config.push_constant_size,
	}

	layout_info := vk.PipelineLayoutCreateInfo {
		sType          = .PIPELINE_LAYOUT_CREATE_INFO,
		setLayoutCount = u32(len(config.descriptor_set_layouts)),
		pSetLayouts    = raw_data(config.descriptor_set_layouts),
	}
	if config.push_constant_size > 0 {
		layout_info.pushConstantRangeCount = 1
		layout_info.pPushConstantRanges = &push_range
	}

	result := vk.CreatePipelineLayout(dev.device, &layout_info, nil, &gp.layout)
	if result != .SUCCESS {
		log.errorf("gpu/pipeline: vkCreatePipelineLayout failed: %v", result)
		return {}, false
	}

	// Create the graphics pipeline
	stage_count: u32 = config.depth_only ? 1 : 2
	pipeline_info := vk.GraphicsPipelineCreateInfo {
		sType               = .GRAPHICS_PIPELINE_CREATE_INFO,
		stageCount          = stage_count,
		pStages             = &stages[0],
		pVertexInputState   = &vertex_input,
		pInputAssemblyState = &input_assembly,
		pViewportState      = &viewport_state,
		pRasterizationState = &rasterization,
		pMultisampleState   = &multisample,
		pDepthStencilState  = &depth_stencil,
		pColorBlendState    = &color_blend,
		pDynamicState       = &dynamic_state,
		layout              = gp.layout,
		renderPass          = config.render_pass,
		subpass             = 0,
	}

	result = vk.CreateGraphicsPipelines(dev.device, 0, 1, &pipeline_info, nil, &gp.pipeline)
	if result != .SUCCESS {
		log.errorf("gpu/pipeline: vkCreateGraphicsPipelines failed: %v", result)
		vk.DestroyPipelineLayout(dev.device, gp.layout, nil)
		return {}, false
	}

	return gp, true
}

destroy_graphics_pipeline :: proc(dev: ^gpu.Gpu_Device, p: ^Graphics_Pipeline) {
	if p.pipeline != 0 {
		vk.DestroyPipeline(dev.device, p.pipeline, nil)
		p.pipeline = 0
	}
	if p.layout != 0 {
		vk.DestroyPipelineLayout(dev.device, p.layout, nil)
		p.layout = 0
	}
}

@(private)
to_vk_blend_factor :: proc(factor: bk.Blend_Factor) -> vk.BlendFactor {
	switch factor {
	case .Zero:
		return .ZERO
	case .One:
		return .ONE
	case .Src_Alpha:
		return .SRC_ALPHA
	case .One_Minus_Src_Alpha:
		return .ONE_MINUS_SRC_ALPHA
	}
	return .ONE
}

@(private)
to_vk_compare_op :: proc(op: bk.Compare_Op) -> vk.CompareOp {
	switch op {
	case .Never:
		return .NEVER
	case .Less:
		return .LESS
	case .Less_Or_Equal:
		return .LESS_OR_EQUAL
	case .Equal:
		return .EQUAL
	case .Greater_Or_Equal:
		return .GREATER_OR_EQUAL
	case .Greater:
		return .GREATER
	case .Not_Equal:
		return .NOT_EQUAL
	case .Always:
		return .ALWAYS
	}
	return .ALWAYS
}

@(private)
to_vk_stencil_op :: proc(op: bk.Stencil_Op) -> vk.StencilOp {
	switch op {
	case .Keep:
		return .KEEP
	case .Zero:
		return .ZERO
	case .Replace:
		return .REPLACE
	case .Increment_Clamp:
		return .INCREMENT_AND_CLAMP
	case .Decrement_Clamp:
		return .DECREMENT_AND_CLAMP
	case .Invert:
		return .INVERT
	}
	return .KEEP
}

@(private)
to_vk_stencil_face :: proc(face: bk.Stencil_Face_State, state: bk.Stencil_State) -> vk.StencilOpState {
	return {
		failOp      = to_vk_stencil_op(face.fail_op),
		passOp      = to_vk_stencil_op(face.pass_op),
		depthFailOp = to_vk_stencil_op(face.depth_fail_op),
		compareOp   = to_vk_compare_op(face.compare_op),
		compareMask = u32(state.read_mask),
		writeMask   = u32(state.write_mask),
		reference   = u32(state.reference),
	}
}