package renderer

// Reference directional-shadow adapter for examples/default app pipelines.
// Do not grow this into a point-light shadow engine; generic Render_IR
// render-target/pass substrate should enable external lighting engines.

import "core:log"
import "core:math"
import "core:mem"
import glsl "core:math/linalg/glsl"
import bk "../backend"
import "../resource"

MAX_SHADOW_CASTERS :: 4
DEFAULT_SHADOW_ATLAS_SIZE :: 4096

// Push constants for shadow depth pass: light_proj_view + model = 128 bytes
Push_Constants_Shadow :: struct {
	light_proj_view: glsl.mat4x4,  // 64 bytes
	model:           glsl.mat4x4,  // 64 bytes
}

// std140-compatible UBO for shadow data passed to lit fragment shader
Shadow_UBO_Data :: struct {
	light_matrices:       [MAX_SHADOW_CASTERS]glsl.mat4x4,  // 256 bytes
	shadow_count:         [4]u32,                            // 16 bytes (x = count)
	shadow_params:        [4]f32,                            // 16 bytes (x = bias, y = normal_bias, z = atlas_size, w = contact_steps)
	shadow_light_indices: [4]u32,                            // 16 bytes (maps shadow map idx -> light idx)
	atlas_regions:        [MAX_SHADOW_CASTERS][4]f32,        // 64 bytes (offset_u, offset_v, scale_u, scale_v)
}

Shadow_State :: struct {
	enabled:            bool,
	bias_constant:      f32,
	bias_slope:         f32,
	shadow_half_extent: f32,
	shadow_depth:       f32,

	// Depth-only render pass
	render_pass:    bk.Render_Pass_Handle,

	// Shadow atlas
	atlas_texture:     bk.Texture_Handle,
	atlas_framebuffer: bk.Framebuffer_Handle,
	atlas_size:        u32,
	sampler:           bk.Sampler_Handle,

	// Depth prepass
	depth_prepass_texture:     bk.Texture_Handle,
	depth_prepass_framebuffer: bk.Framebuffer_Handle,
	depth_prepass_sampler:     bk.Sampler_Handle,
	depth_prepass_width:       u32,
	depth_prepass_height:      u32,
	depth_format:              bk.Format,

	// Light-space matrices
	light_matrices:       [MAX_SHADOW_CASTERS]glsl.mat4x4,
	shadow_light_indices: [MAX_SHADOW_CASTERS]u32,
	shadow_count:         u32,

	// Shadow UBO
	ubo_buffers: [bk.MAX_FRAMES_IN_FLIGHT]bk.Buffer_Handle,

	// Descriptor set for shadow data (set=2 in lit pipeline)
	descriptor_pool:       bk.Descriptor_Handle,
	descriptor_set_layout: bk.Descriptor_Handle,
	descriptor_sets:       [bk.MAX_FRAMES_IN_FLIGHT]bk.Descriptor_Handle,

	// Shadow depth pipelines (one per layout variant)
	depth_pipelines: [bk.MAX_LAYOUT_VARIANTS]bk.Pipeline_Handle,

	initialized: bool,
}

init_shadow_state :: proc(state: ^Shadow_State, b: ^bk.Backend, screen_width, screen_height: u32) -> bool {
	state.atlas_size = DEFAULT_SHADOW_ATLAS_SIZE
	state.bias_constant = 1.25
	state.bias_slope = 1.75
	state.shadow_half_extent = 15.0
	state.shadow_depth = 40.0
	state.depth_format = b.get_depth_format()

	// Create depth-only render pass (shared by atlas + depth prepass)
	rp, rp_ok := b.create_render_pass(bk.Render_Pass_Desc{
		has_depth    = true,
		depth_format = state.depth_format,
		depth_only   = true,
		depth_load_op = .Clear,
		depth_store_op = .Store,
		depth_final_layout = .Depth_Stencil_Read_Only,
	})
	if !rp_ok {
		log.error("gpu/renderer: failed to create shadow render pass")
		return false
	}
	state.render_pass = rp

	// Create shadow atlas depth image
	atlas_tex, atlas_ok := b.create_image(bk.Texture_Desc{
		width  = state.atlas_size,
		height = state.atlas_size,
		format = state.depth_format,
		usage  = {.Depth_Stencil_Attachment, .Sampled},
	})
	if !atlas_ok {
		log.error("gpu/renderer: failed to create shadow atlas image")
		shutdown_shadow_state(state, b)
		return false
	}

	if !b.create_image_view(atlas_tex, state.depth_format, {.Depth}) {
		b.destroy_image(atlas_tex)
		log.error("gpu/renderer: failed to create shadow atlas image view")
		shutdown_shadow_state(state, b)
		return false
	}
	state.atlas_texture = atlas_tex

	// Create framebuffer for the atlas
	fb, fb_ok := b.create_framebuffer({
		pass = state.render_pass,
		depth_view = atlas_tex,
		width = state.atlas_size,
		height = state.atlas_size,
		layers = 1,
	})
	if !fb_ok {
		log.error("gpu/renderer: failed to create shadow atlas framebuffer")
		shutdown_shadow_state(state, b)
		return false
	}
	state.atlas_framebuffer = fb

	// Create comparison sampler for hardware shadow-map depth tests.
	sampler, sampler_ok := b.create_sampler(bk.Sampler_Desc{
		mag_filter     = .Linear,
		min_filter     = .Linear,
		address_mode_u = .Clamp_To_Border,
		address_mode_v = .Clamp_To_Border,
		enable_compare = true,
		compare_op     = .Less_Or_Equal,
	})
	if !sampler_ok {
		log.error("gpu/renderer: failed to create shadow sampler")
		shutdown_shadow_state(state, b)
		return false
	}
	state.sampler = sampler

	// Create depth prepass resources
	if !create_depth_prepass(state, b, screen_width, screen_height) {
		log.error("gpu/renderer: failed to create depth prepass")
		shutdown_shadow_state(state, b)
		return false
	}

	// Create regular sampler for depth prepass (reads depth values, not comparison)
	dp_sampler, dp_sampler_ok := b.create_sampler(bk.Sampler_Desc{
		mag_filter     = .Nearest,
		min_filter     = .Nearest,
		address_mode_u = .Clamp_To_Edge,
		address_mode_v = .Clamp_To_Edge,
	})
	if !dp_sampler_ok {
		log.error("gpu/renderer: failed to create depth prepass sampler")
		shutdown_shadow_state(state, b)
		return false
	}
	state.depth_prepass_sampler = dp_sampler

	// Create descriptor set layout: binding 0 = UBO, binding 1 = shadow atlas, binding 3 = depth prepass
	// Note: binding 2 is skipped because Luma splits sampler2DShadow at binding 1
	// into texture(1) + sampler(2), consuming binding 2 internally.
	layout_bindings := [3]bk.Descriptor_Set_Layout_Binding{
		{binding = 0, type = .Uniform_Buffer,          count = 1, stages = {.Fragment}},
		{binding = 1, type = .Combined_Image_Sampler,  count = 1, stages = {.Fragment}},
		{binding = 3, type = .Combined_Image_Sampler,  count = 1, stages = {.Fragment}},
	}
	layout, layout_ok := b.create_descriptor_set_layout(layout_bindings[:])
	if !layout_ok {
		log.error("gpu/renderer: failed to create shadow descriptor set layout")
		shutdown_shadow_state(state, b)
		return false
	}
	state.descriptor_set_layout = layout

	// Create descriptor pool
	types := [2]bk.Descriptor_Type{.Uniform_Buffer, .Combined_Image_Sampler}
	counts := [2]u32{bk.MAX_FRAMES_IN_FLIGHT, bk.MAX_FRAMES_IN_FLIGHT * 2}
	pool, pool_ok := b.create_descriptor_pool(bk.MAX_FRAMES_IN_FLIGHT, types[:], counts[:])
	if !pool_ok {
		log.error("gpu/renderer: failed to create shadow descriptor pool")
		shutdown_shadow_state(state, b)
		return false
	}
	state.descriptor_pool = pool

	// Create per-frame UBO buffers and descriptor sets
	for i in 0..<bk.MAX_FRAMES_IN_FLIGHT {
		buf, buf_ok := b.create_buffer(bk.Buffer_Desc{
			size   = u64(size_of(Shadow_UBO_Data)),
			usage  = {.Uniform},
			memory = {.Host_Visible, .Host_Coherent},
		})
		if !buf_ok {
			log.errorf("gpu/renderer: failed to create shadow UBO buffer %d", i)
			shutdown_shadow_state(state, b)
			return false
		}
		state.ubo_buffers[i] = buf

		// Persistently map
		mapped := b.map_buffer(buf)
		if mapped == nil {
			log.errorf("gpu/renderer: failed to map shadow UBO buffer %d", i)
			shutdown_shadow_state(state, b)
			return false
		}

		// Allocate descriptor set
		ds, ds_ok := b.allocate_descriptor_set(state.descriptor_pool, state.descriptor_set_layout)
		if !ds_ok {
			log.errorf("gpu/renderer: failed to allocate shadow descriptor set %d", i)
			shutdown_shadow_state(state, b)
			return false
		}
		state.descriptor_sets[i] = ds

		// Write UBO to binding 0
		b.update_descriptor_buffer(ds, 0, buf, u64(size_of(Shadow_UBO_Data)))

		// Write shadow atlas to binding 1
		b.update_descriptor_image(ds, 1, state.atlas_texture, state.sampler, .Depth_Stencil_Read_Only)

		// Write depth prepass to binding 3 (binding 2 reserved by Luma sampler split)
		b.update_descriptor_image(ds, 3, state.depth_prepass_texture, state.depth_prepass_sampler, .Depth_Stencil_Read_Only)
	}

	state.initialized = true
	log.info("gpu/renderer: shadow system initialized (atlas + contact shadows)")
	return true
}

shutdown_shadow_state :: proc(state: ^Shadow_State, b: ^bk.Backend) {
	b.wait_idle()

	for v in 0..<bk.MAX_LAYOUT_VARIANTS {
		b.destroy_graphics_pipeline(state.depth_pipelines[v])
	}

	for i in 0..<bk.MAX_FRAMES_IN_FLIGHT {
		b.unmap_buffer(state.ubo_buffers[i])
		b.destroy_buffer(state.ubo_buffers[i])
	}

	b.destroy_descriptor_pool(state.descriptor_pool)
	b.destroy_descriptor_set_layout(state.descriptor_set_layout)

	b.destroy_sampler(state.sampler)
	b.destroy_sampler(state.depth_prepass_sampler)

	destroy_depth_prepass(state, b)

	b.destroy_framebuffer(state.atlas_framebuffer)
	b.destroy_image(state.atlas_texture)

	b.destroy_render_pass(state.render_pass)

	state.initialized = false
}

// Recreate depth prepass on window resize. Call after swapchain recreation.
on_shadow_resize :: proc(state: ^Shadow_State, b: ^bk.Backend, width, height: u32) {
	if !state.initialized { return }
	if width == state.depth_prepass_width && height == state.depth_prepass_height { return }

	b.wait_idle()

	destroy_depth_prepass(state, b)
	if !create_depth_prepass(state, b, width, height) {
		log.error("gpu/renderer: failed to recreate depth prepass on resize")
		return
	}

	// Update descriptor sets with new depth prepass image
	for i in 0..<bk.MAX_FRAMES_IN_FLIGHT {
		b.update_descriptor_image(state.descriptor_sets[i], 3, state.depth_prepass_texture, state.depth_prepass_sampler, .Depth_Stencil_Read_Only)
	}
}

// Execute shadow atlas pass + depth prepass for contact shadows.
// Must be called BEFORE ensure_render_pass (before main render pass begins).
execute_shadow_pass :: proc(
	state: ^Shadow_State,
	b: ^bk.Backend,
	ctx: bk.Frame_Context,
	q: ^Draw_Queue_3D,
	light_state: ^Light_State,
	camera_target: glsl.vec3,
	camera_proj_view: glsl.mat4x4,
	frame_index: u32,
) {
	if !state.initialized || !state.enabled || q.count == 0 {
		state.shadow_count = 0
		// Transition atlas and depth prepass to valid layout
		transition_atlas_empty(state, b, ctx)
		transition_depth_prepass_empty(state, b, ctx)
		return
	}

	// --- Phase 1: Shadow atlas (directional lights) ---

	shadow_idx: u32 = 0
	for i in 0..<light_state.light_count {
		if shadow_idx >= MAX_SHADOW_CASTERS { break }
		light := &light_state.lights[i]
		if !light.enabled || !light.casts_shadow { continue }
		// Only directional lights get shadow maps
		if light.type != 0 { continue }

		light_dir := glsl.vec3{light.position.x, light.position.y, light.position.z}
		light_mat := compute_light_matrix(light_dir, camera_target, state.shadow_half_extent, state.shadow_depth)
		state.light_matrices[shadow_idx] = light_mat
		state.shadow_light_indices[shadow_idx] = u32(i)
		shadow_idx += 1
	}

	state.shadow_count = shadow_idx

	if shadow_idx > 0 {
		// Render directional shadow atlas
		b.begin_render_pass(ctx, {
			pass = state.render_pass,
			framebuffer = state.atlas_framebuffer,
			width = state.atlas_size,
			height = state.atlas_size,
			clear_depth = 1,
		})
		b.set_depth_bias(ctx, state.bias_constant, state.bias_slope)

		for i in 0..<shadow_idx {
			region := compute_atlas_region(i, shadow_idx, state.atlas_size)

			b.set_viewport(ctx, f32(region.x), f32(region.y), f32(region.w), f32(region.h))
			b.set_scissor(ctx, region.x, region.y, region.w, region.h)

			render_shadow_meshes(q, b, ctx, state.depth_pipelines, state.light_matrices[i])
		}

		b.end_render_pass(ctx)
	} else {
		transition_atlas_empty(state, b, ctx)
	}

	// --- Phase 2: Depth prepass (camera view, for contact shadows) ---

	execute_depth_prepass(state, b, ctx, q, camera_proj_view)

	// --- Update shadow UBO ---
	update_shadow_ubo(state, b, frame_index)
}

// --- Depth prepass ---

@(private = "file")
create_depth_prepass :: proc(state: ^Shadow_State, b: ^bk.Backend, width, height: u32) -> bool {
	tex, tex_ok := b.create_image(bk.Texture_Desc{
		width  = width,
		height = height,
		format = state.depth_format,
		usage  = {.Depth_Stencil_Attachment, .Sampled},
	})
	if !tex_ok { return false }

	if !b.create_image_view(tex, state.depth_format, {.Depth}) {
		b.destroy_image(tex)
		return false
	}
	state.depth_prepass_texture = tex

	fb, fb_ok := b.create_framebuffer({
		pass = state.render_pass,
		depth_view = tex,
		width = width,
		height = height,
		layers = 1,
	})
	if !fb_ok {
		b.destroy_image(tex)
		state.depth_prepass_texture = {}
		return false
	}
	state.depth_prepass_framebuffer = fb
	state.depth_prepass_width = width
	state.depth_prepass_height = height
	return true
}

@(private = "file")
destroy_depth_prepass :: proc(state: ^Shadow_State, b: ^bk.Backend) {
	b.destroy_framebuffer(state.depth_prepass_framebuffer)
	b.destroy_image(state.depth_prepass_texture)
	state.depth_prepass_framebuffer = {}
	state.depth_prepass_texture = {}
	state.depth_prepass_width = 0
	state.depth_prepass_height = 0
}

@(private = "file")
execute_depth_prepass :: proc(state: ^Shadow_State, b: ^bk.Backend, ctx: bk.Frame_Context, q: ^Draw_Queue_3D, proj_view: glsl.mat4x4) {
	if state.depth_prepass_width == 0 || state.depth_prepass_height == 0 {
		return
	}

	b.begin_render_pass(ctx, {
		pass = state.render_pass,
		framebuffer = state.depth_prepass_framebuffer,
		width = state.depth_prepass_width,
		height = state.depth_prepass_height,
		clear_depth = 1,
	})

	b.set_viewport(ctx, 0, 0, f32(state.depth_prepass_width), f32(state.depth_prepass_height))
	b.set_scissor(ctx, 0, 0, state.depth_prepass_width, state.depth_prepass_height)

	b.set_depth_bias(ctx, 0, 0)  // no bias for depth prepass

	// Render all meshes with camera proj_view
	render_shadow_meshes(q, b, ctx, state.depth_pipelines, proj_view)

	b.end_render_pass(ctx)
}

// --- Atlas region computation ---

Atlas_Region :: struct {
	x, y: i32,
	w, h: u32,
}

// Compute pixel-space atlas region for a shadow caster.
// With 1 caster: full atlas. With 2-4: subdivide into quadrants.
compute_atlas_region :: proc(index, count, atlas_size: u32) -> Atlas_Region {
	if count <= 1 {
		return {0, 0, atlas_size, atlas_size}
	}
	half := atlas_size / 2
	switch index {
	case 0: return {0, 0, half, half}
	case 1: return {i32(half), 0, half, half}
	case 2: return {0, i32(half), half, half}
	case 3: return {i32(half), i32(half), half, half}
	}
	return {0, 0, half, half}
}

// Compute normalized UV region for shader sampling.
compute_atlas_uv_region :: proc(index, count: u32) -> [4]f32 {
	if count <= 1 {
		return {0, 0, 1, 1}
	}
	switch index {
	case 0: return {0, 0, 0.5, 0.5}
	case 1: return {0.5, 0, 0.5, 0.5}
	case 2: return {0, 0.5, 0.5, 0.5}
	case 3: return {0.5, 0.5, 0.5, 0.5}
	}
	return {0, 0, 0.5, 0.5}
}

// --- Internal helpers ---

@(private = "file")
update_shadow_ubo :: proc(state: ^Shadow_State, b: ^bk.Backend, frame_index: u32) {
	ubo: Shadow_UBO_Data
	for i in 0..<state.shadow_count {
		ubo.light_matrices[i] = state.light_matrices[i]
		ubo.shadow_light_indices[i] = state.shadow_light_indices[i]
		ubo.atlas_regions[i] = compute_atlas_uv_region(i, state.shadow_count)
	}
	ubo.shadow_count = {state.shadow_count, 0, 0, 0}
	ubo.shadow_params = {0.003, 0.03, f32(state.atlas_size), 6}  // x = depth bias, y = normal bias, w = contact shadow steps

	mapped := b.get_buffer_mapped(state.ubo_buffers[frame_index])
	if mapped != nil {
		mem.copy(mapped, &ubo, size_of(Shadow_UBO_Data))
	}
}

// Render all meshes from the draw queue into a depth buffer.
// Does NOT consume the queue (does not reset count).
// Switches pipeline per-draw based on the mesh's layout variant.
@(private = "file")
render_shadow_meshes :: proc(
	q: ^Draw_Queue_3D,
	b: ^bk.Backend,
	ctx: bk.Frame_Context,
	depth_pipelines: [bk.MAX_LAYOUT_VARIANTS]bk.Pipeline_Handle,
	light_proj_view: glsl.mat4x4,
) {
	current_variant: bk.Layout_Variant = max(bk.Layout_Variant)
	current_pipeline: bk.Pipeline_Handle

	for i in 0..<q.count {
		draw := &q.commands[i]

		if draw.layout_variant != current_variant {
			current_variant = draw.layout_variant
			current_pipeline = depth_pipelines[current_variant]
			if current_pipeline == bk.NULL_PIPELINE {
				continue
			}
			b.bind_graphics_pipeline(ctx, current_pipeline)
		}
		if current_pipeline == bk.NULL_PIPELINE {
			continue
		}

		pc := Push_Constants_Shadow{
			light_proj_view = light_proj_view,
			model           = draw.model,
		}
		b.push_constants(ctx, current_pipeline, {.Vertex}, 0, size_of(Push_Constants_Shadow), &pc)

		vb, ib, index_count, mesh_ok := resource.get_mesh_buffers(q.resource_state, draw.mesh_id)
		if !mesh_ok { continue }

		b.bind_vertex_buffer(ctx, vb)
		b.bind_index_buffer(ctx, ib)
		b.draw_indexed(ctx, u32(index_count), 1, 0, 0, 0)
	}
}

// Compute orthographic light-space matrix for a directional light.
compute_light_matrix :: proc(light_dir, camera_target: glsl.vec3, half_extent, depth: f32) -> glsl.mat4x4 {
	dir := normalize(light_dir)
	light_pos := camera_target - dir * (depth * 0.5)

	world_up: glsl.vec3
	if math.abs(dir.y) > 0.99 {
		world_up = {0, 0, 1}
	} else {
		world_up = {0, 1, 0}
	}
	view := look_at(light_pos, camera_target, world_up)
	proj := ortho_shadow(-half_extent, half_extent, -half_extent, half_extent, 0, depth)
	return proj * view
}

@(private = "file")
ortho_shadow :: proc(left, right, bottom, top, near, far: f32) -> glsl.mat4x4 {
	m: glsl.mat4x4
	m[0, 0] = 2.0 / (right - left)
	m[1, 1] = -2.0 / (top - bottom)
	m[2, 2] = -1.0 / (far - near)
	m[0, 3] = -(right + left) / (right - left)
	m[1, 3] = -(top + bottom) / (top - bottom)
	m[2, 3] = -near / (far - near)
	m[3, 3] = 1.0
	return m
}

@(private = "file")
transition_atlas_empty :: proc(state: ^Shadow_State, b: ^bk.Backend, ctx: bk.Frame_Context) {
	b.begin_render_pass(ctx, {
		pass = state.render_pass,
		framebuffer = state.atlas_framebuffer,
		width = state.atlas_size,
		height = state.atlas_size,
		clear_depth = 1,
	})
	b.end_render_pass(ctx)
}

@(private = "file")
transition_depth_prepass_empty :: proc(state: ^Shadow_State, b: ^bk.Backend, ctx: bk.Frame_Context) {
	if state.depth_prepass_width == 0 || state.depth_prepass_height == 0 { return }

	b.begin_render_pass(ctx, {
		pass = state.render_pass,
		framebuffer = state.depth_prepass_framebuffer,
		width = state.depth_prepass_width,
		height = state.depth_prepass_height,
		clear_depth = 1,
	})
	b.end_render_pass(ctx)
}