package renderer

// Reference lighting adapter for examples/default app pipelines.
// Keep policy minimal here; external lighting engines should own clustering,
// BRDF choices, point-shadow budgets, and shadow-map layout decisions.

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

MAX_LIGHTS :: 336

// std140-compatible GPU light (48 bytes)
GPU_Light :: struct {
	position: [4]f32,  // xyz = pos/dir, w = type (0=dir, 1=point)
	color:    [4]f32,  // rgb = color*intensity, w = enabled (0/1)
	params:   [4]f32,  // x = radius, y = intensity, zw = reserved
}

// std140-compatible UBO layout (16240 bytes, fits in Vulkan minimum 16384)
Light_UBO_Data :: struct {
	proj_view:     glsl.mat4x4,        // 64 bytes, offset 0
	camera_pos:    [4]f32,             // 16 bytes, offset 64
	ambient_color: [4]f32,             // 16 bytes, offset 80
	light_count:   [4]f32,            // 16 bytes, offset 96 (x = count as float)
	lights:        [MAX_LIGHTS]GPU_Light, // 16128 bytes, offset 112
}
// Total: 16240 bytes (336 lights × 48 bytes + 112 byte header)

// Public-facing light data (mirrors gpu.Light)
Light_Data :: struct {
	type:         u32,    // 0 = directional, 1 = point
	enabled:      bool,
	position:     [3]f32,
	color:        [4]f32,
	intensity:    f32,
	radius:       f32,
	casts_shadow: bool,
}

Light_State :: struct {
	lights:        [MAX_LIGHTS]Light_Data,
	light_count:   u32,
	ambient_color: [4]f32,

	// Per-frame UBOs (double-buffered)
	ubo_buffers:         [bk.MAX_FRAMES_IN_FLIGHT]bk.Buffer_Handle,
	ubo_descriptor_sets: [bk.MAX_FRAMES_IN_FLIGHT]bk.Descriptor_Handle,
	ubo_set_layout:      bk.Descriptor_Handle,
	ubo_pool:            bk.Descriptor_Handle,
}

init_light_state :: proc(state: ^Light_State, b: ^bk.Backend) -> bool {
	// Create UBO descriptor set layout
	layout, layout_ok := b.create_descriptor_set_layout({{
		binding = 0,
		type    = .Uniform_Buffer,
		count   = 1,
		stages  = {.Vertex, .Fragment},
	}})
	if !layout_ok {
		log.error("gpu/renderer: failed to create light UBO descriptor set layout")
		return false
	}
	state.ubo_set_layout = layout

	// Create UBO descriptor pool
	types := [1]bk.Descriptor_Type{.Uniform_Buffer}
	counts := [1]u32{bk.MAX_FRAMES_IN_FLIGHT}
	pool, pool_ok := b.create_descriptor_pool(bk.MAX_FRAMES_IN_FLIGHT, types[:], counts[:])
	if !pool_ok {
		b.destroy_descriptor_set_layout(state.ubo_set_layout)
		log.error("gpu/renderer: failed to create light UBO descriptor pool")
		return false
	}
	state.ubo_pool = pool

	// Create per-frame UBO buffers and descriptor sets
	ubo_size := u64(size_of(Light_UBO_Data))
	for i in 0..<bk.MAX_FRAMES_IN_FLIGHT {
		buf, buf_ok := b.create_buffer(bk.Buffer_Desc{
			size   = ubo_size,
			usage  = {.Uniform},
			memory = {.Host_Visible, .Host_Coherent},
		})
		if !buf_ok {
			log.errorf("gpu/renderer: failed to create light UBO buffer %d", i)
			shutdown_light_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 light UBO buffer %d", i)
			shutdown_light_state(state, b)
			return false
		}

		// Allocate descriptor set
		ds, ds_ok := b.allocate_descriptor_set(state.ubo_pool, state.ubo_set_layout)
		if !ds_ok {
			log.errorf("gpu/renderer: failed to allocate light UBO descriptor set %d", i)
			shutdown_light_state(state, b)
			return false
		}
		state.ubo_descriptor_sets[i] = ds

		// Write buffer to descriptor set
		b.update_descriptor_buffer(ds, 0, buf, ubo_size)
	}

	// Default ambient
	state.ambient_color = {0.1, 0.1, 0.1, 1.0}

	return true
}

shutdown_light_state :: proc(state: ^Light_State, b: ^bk.Backend) {
	b.wait_idle()

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

	if state.ubo_pool != bk.NULL_DESCRIPTOR {
		b.destroy_descriptor_pool(state.ubo_pool)
		state.ubo_pool = bk.NULL_DESCRIPTOR
	}
	if state.ubo_set_layout != bk.NULL_DESCRIPTOR {
		b.destroy_descriptor_set_layout(state.ubo_set_layout)
		state.ubo_set_layout = bk.NULL_DESCRIPTOR
	}
}

// Update the light UBO for the current frame
update_light_ubo :: proc(state: ^Light_State, b: ^bk.Backend, frame_index: u32, proj_view: glsl.mat4x4, camera_pos: [3]f32) {
	ubo: Light_UBO_Data
	ubo.proj_view = proj_view
	ubo.camera_pos = {camera_pos.x, camera_pos.y, camera_pos.z, 0}
	ubo.ambient_color = state.ambient_color
	ubo.light_count = {f32(state.light_count), 0, 0, 0}

	for i in 0..<state.light_count {
		light := &state.lights[i]
		ubo.lights[i] = GPU_Light{
			position = {light.position.x, light.position.y, light.position.z, f32(light.type)},
			color    = {light.color.x * light.intensity, light.color.y * light.intensity, light.color.z * light.intensity, light.enabled ? 1.0 : 0.0},
			params   = {light.radius, light.intensity, 0, 0},
		}
	}

	// Copy to mapped buffer
	mapped := b.get_buffer_mapped(state.ubo_buffers[frame_index])
	if mapped != nil {
		mem.copy(mapped, &ubo, size_of(Light_UBO_Data))
	} else {
		log.error("gpu/renderer: light UBO buffer not mapped!")
	}
}

// Light management
add_light :: proc(state: ^Light_State, light: Light_Data) -> i32 {
	if state.light_count >= MAX_LIGHTS {
		log.warnf("gpu/renderer: light limit reached (%d), light dropped", MAX_LIGHTS)
		return -1
	}
	idx := state.light_count
	state.lights[idx] = light
	state.light_count += 1
	return i32(idx)
}

set_light :: proc(state: ^Light_State, index: i32, light: Light_Data) {
	if index < 0 || u32(index) >= state.light_count {
		return
	}
	state.lights[index] = light
}

remove_light :: proc(state: ^Light_State, index: i32) {
	if index < 0 || u32(index) >= state.light_count {
		return
	}
	// Shift remaining lights down
	for i in u32(index)..<state.light_count - 1 {
		state.lights[i] = state.lights[i + 1]
	}
	state.light_count -= 1
}

clear_lights :: proc(state: ^Light_State) {
	state.light_count = 0
}