package resource

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

MAX_TEXTURES         :: 256
WHITE_TEXTURE_ID     :: 0
NORMAL_TEXTURE_ID    :: 1
MAX_MODELS           :: 64
MAX_COMPUTE_SHADERS  :: 32
MAX_STORAGE_BUFFERS  :: 64

Internal_Texture :: struct {
	texture:        bk.Texture_Handle,
	sampler:        bk.Sampler_Handle,
	descriptor_set: bk.Descriptor_Handle,
	width:          i32,
	height:         i32,
	owns_texture:   bool,
	active:         bool,
}

Internal_Model :: struct {
	mesh_ids:              []u32,
	texture_ids:           []u32,
	normal_map_ids:        []u32,
	material_colors:       [][4]f32,
	material_double_sided: []bool,
	material_metallic:     []f32,
	material_roughness:    []f32,
	material_emissive:     [][4]f32,
	mesh_count:            int,
	transform:             glsl.mat4x4,
	active:                bool,
}

Internal_Compute_Shader :: struct {
	shader:    bk.Shader_Handle,
	pipeline:  bk.Pipeline_Handle,
	desc_layout: bk.Descriptor_Handle,
	desc_pool:   bk.Descriptor_Handle,
	desc_sets: [bk.MAX_FRAMES_IN_FLIGHT]bk.Descriptor_Handle,
	active:    bool,
}

Internal_Storage_Buffer :: struct {
	buffer: bk.Buffer_Handle,
	size:   int,
	active: bool,
}

Resource_State :: struct {
	textures:         [MAX_TEXTURES]Internal_Texture,
	meshes:           [MAX_MESHES]Internal_Mesh,
	models:           [MAX_MODELS]Internal_Model,
	compute_shaders:  [MAX_COMPUTE_SHADERS]Internal_Compute_Shader,
	storage_buffers:  [MAX_STORAGE_BUFFERS]Internal_Storage_Buffer,
	desc_layout:      bk.Descriptor_Handle,
	desc_pool:        bk.Descriptor_Handle,
	backend:          ^bk.Backend,
}

init_resources :: proc(b: ^bk.Backend) -> (state: Resource_State, ok: bool) {
	state.backend = b

	// Create descriptor set layout (single combined image sampler at binding 0)
	layout, layout_ok := b.create_descriptor_set_layout({{
		binding = 0,
		type    = .Combined_Image_Sampler,
		count   = 1,
		stages  = {.Fragment},
	}})
	if !layout_ok {
		return {}, false
	}
	state.desc_layout = layout

	// Create descriptor pool
	types := [1]bk.Descriptor_Type{.Combined_Image_Sampler}
	counts := [1]u32{MAX_TEXTURES + 1}
	pool, pool_ok := b.create_descriptor_pool(MAX_TEXTURES + 1, types[:], counts[:])
	if !pool_ok {
		b.destroy_descriptor_set_layout(state.desc_layout)
		return {}, false
	}
	state.desc_pool = pool

	// Create 1x1 white default texture in slot 0
	white_ok := create_white_texture(&state, b)
	if !white_ok {
		log.error("gpu/resource: failed to create white default texture")
		b.destroy_descriptor_pool(state.desc_pool)
		b.destroy_descriptor_set_layout(state.desc_layout)
		return {}, false
	}

	// Create 1x1 flat-normal default texture in slot 1 (tangent-space up: 0,0,1)
	normal_ok := create_normal_texture(&state, b)
	if !normal_ok {
		log.error("gpu/resource: failed to create normal default texture")
		b.destroy_descriptor_pool(state.desc_pool)
		b.destroy_descriptor_set_layout(state.desc_layout)
		return {}, false
	}

	return state, true
}

allocate_model_slot :: proc(state: ^Resource_State) -> (u32, bool) {
	for i in 1..<u32(MAX_MODELS) {
		if !state.models[i].active {
			return i, true
		}
	}
	log.error("gpu/resource: model pool exhausted")
	return 0, false
}

shutdown_models :: proc(state: ^Resource_State, b: ^bk.Backend) {
	for i in 0..<u32(MAX_MODELS) {
		if state.models[i].active {
			unload_model_resource(state, b, i)
		}
	}
}

shutdown_resources :: proc(state: ^Resource_State, b: ^bk.Backend) {
	// Destroy compute shaders and storage buffers
	shutdown_compute(state, b)

	// Destroy models (which unload their meshes/textures)
	shutdown_models(state, b)

	// Destroy remaining meshes
	shutdown_meshes(state, b)

	// Destroy descriptor pool first (frees all descriptor sets)
	b.destroy_descriptor_pool(state.desc_pool)

	// Now safe to destroy textures (samplers, images)
	for i in 0..<MAX_TEXTURES {
		if state.textures[i].active {
			destroy_texture_slot(b, &state.textures[i])
		}
	}

	b.destroy_descriptor_set_layout(state.desc_layout)
}

get_descriptor_set_layout :: proc(state: ^Resource_State) -> bk.Descriptor_Handle {
	return state.desc_layout
}

get_texture_descriptor_set :: proc(state: ^Resource_State, id: u32) -> bk.Descriptor_Handle {
	if id < MAX_TEXTURES && state.textures[id].active {
		return state.textures[id].descriptor_set
	}
	return state.textures[WHITE_TEXTURE_ID].descriptor_set
}

import_texture_view :: proc(
	state: ^Resource_State,
	b: ^bk.Backend,
	texture: bk.Texture_Handle,
	width, height: u32,
) -> (id: u32, ok: bool) {
	if state == nil || b == nil || texture == bk.NULL_TEXTURE || width == 0 || height == 0 {
		return 0, false
	}
	slot, slot_ok := allocate_texture_slot(state)
	if !slot_ok {
		log.error("gpu/resource: texture pool full")
		return 0, false
	}

	tex := &state.textures[slot]
	sam_handle, sam_ok := b.create_sampler(bk.Sampler_Desc{
		mag_filter     = .Linear,
		min_filter     = .Linear,
		address_mode_u = .Clamp_To_Edge,
		address_mode_v = .Clamp_To_Edge,
		mipmap_mode    = .Linear,
	})
	if !sam_ok {
		return 0, false
	}

	desc_handle, desc_ok := b.allocate_descriptor_set(state.desc_pool, state.desc_layout)
	if !desc_ok {
		b.destroy_sampler(sam_handle)
		return 0, false
	}

	b.update_descriptor_image(desc_handle, 0, texture, sam_handle, .Shader_Read_Only)
	tex^ = {
		texture = texture,
		sampler = sam_handle,
		descriptor_set = desc_handle,
		width = i32(width),
		height = i32(height),
		owns_texture = false,
		active = true,
	}
	return slot, true
}

@(private)
create_white_texture :: proc(state: ^Resource_State, b: ^bk.Backend) -> bool {
	white_pixel := [4]u8{255, 255, 255, 255}

	id, ok := load_texture_from_pixels(
		state, b,
		raw_data(&white_pixel), 1, 1,
	)
	if !ok {
		return false
	}

	assert(id == WHITE_TEXTURE_ID)
	return true
}

@(private)
create_normal_texture :: proc(state: ^Resource_State, b: ^bk.Backend) -> bool {
	// Flat normal in tangent space: (0,0,1) encoded as (128,128,255)
	normal_pixel := [4]u8{128, 128, 255, 255}

	id, ok := load_texture_from_pixels(
		state, b,
		raw_data(&normal_pixel), 1, 1,
	)
	if !ok {
		return false
	}

	assert(id == NORMAL_TEXTURE_ID)
	return true
}

// --- Compute Shader Resources ---

create_compute_shader_resource :: proc(
	state: ^Resource_State,
	b: ^bk.Backend,
	name: string,
	data: []u8,
	format: bk.Shader_Format,
	num_buffers: u32,
	push_constant_size: u32 = 0,
) -> (id: u32, ok: bool) {
	if len(data) == 0 {
		log.errorf("gpu/resource: compute shader '%s' has no bytecode", name)
		return 0, false
	}
	if format != bk.REQUIRED_SHADER_FORMAT {
		log.errorf("gpu/resource: compute shader '%s' format does not match backend requirement", name)
		return 0, false
	}
	if !bk.supports_push_constant_size(b.capabilities, push_constant_size) {
		log.errorf(
			"gpu/resource: compute shader '%s' push constant size %d exceeds backend implemented limit %d",
			name,
			push_constant_size,
			b.capabilities.max_push_constant_size,
		)
		return 0, false
	}

	slot: u32
	slot_found := false
	for i in 1..<u32(MAX_COMPUTE_SHADERS) {
		if !state.compute_shaders[i].active {
			slot = i
			slot_found = true
			break
		}
	}
	if !slot_found {
		log.error("gpu/resource: compute shader pool exhausted")
		return 0, false
	}

	shader_handle, shader_ok := b.create_shader_module(bk.Shader_Module_Desc{
		stage  = .Compute,
		format = format,
		name   = name,
		entry  = "main",
		data   = data,
	})
	if !shader_ok {
		return 0, false
	}

	pip_handle, pip_ok := b.create_compute_pipeline(shader_handle, num_buffers, push_constant_size)
	if !pip_ok {
		b.destroy_shader(shader_handle)
		return 0, false
	}

	cs := &state.compute_shaders[slot]
	cs.shader = shader_handle
	cs.pipeline = pip_handle

	if num_buffers > 0 {
		bindings := make([]bk.Descriptor_Set_Layout_Binding, num_buffers, context.temp_allocator)
		for i in 0..<num_buffers {
			bindings[i] = {
				binding = i,
				type    = .Storage_Buffer,
				count   = 1,
				stages  = {.Compute},
			}
		}

		desc_layout, desc_layout_ok := b.create_descriptor_set_layout(bindings)
		if !desc_layout_ok {
			b.destroy_compute_pipeline(pip_handle)
			b.destroy_shader(shader_handle)
			cs^ = {}
			return 0, false
		}
		cs.desc_layout = desc_layout

		types := [1]bk.Descriptor_Type{.Storage_Buffer}
		counts := [1]u32{num_buffers * bk.MAX_FRAMES_IN_FLIGHT}
		desc_pool, desc_pool_ok := b.create_descriptor_pool(bk.MAX_FRAMES_IN_FLIGHT, types[:], counts[:])
		if !desc_pool_ok {
			b.destroy_descriptor_set_layout(cs.desc_layout)
			b.destroy_compute_pipeline(pip_handle)
			b.destroy_shader(shader_handle)
			cs^ = {}
			return 0, false
		}
		cs.desc_pool = desc_pool

		for frame in 0..<bk.MAX_FRAMES_IN_FLIGHT {
			desc_set, desc_set_ok := b.allocate_descriptor_set(cs.desc_pool, cs.desc_layout)
			if !desc_set_ok {
				b.destroy_descriptor_pool(cs.desc_pool)
				b.destroy_descriptor_set_layout(cs.desc_layout)
				b.destroy_compute_pipeline(pip_handle)
				b.destroy_shader(shader_handle)
				cs^ = {}
				return 0, false
			}
			cs.desc_sets[frame] = desc_set
		}
	}

	cs.active = true
	return slot, true
}

load_compute_shader_resource :: proc(
	state: ^Resource_State,
	b: ^bk.Backend,
	path: cstring,
	num_buffers: u32,
	push_constant_size: u32 = 0,
) -> (id: u32, ok: bool) {
	// Read shader file
	data, read_err := os.read_entire_file(string(path), context.allocator)
	if read_err != nil {
		log.errorf("gpu/resource: failed to read compute shader: %s", path)
		return 0, false
	}
	defer delete(data)

	format: bk.Shader_Format
	if strings.has_suffix(string(path), ".spv") {
		format = .SPIRV
	} else if strings.has_suffix(string(path), ".hlsl") {
		format = .HLSL
	} else {
		log.error("gpu/resource: compute shader files must be .spv or .hlsl")
		return 0, false
	}
	return create_compute_shader_resource(state, b, string(path), data, format, num_buffers, push_constant_size)
}

unload_compute_shader_resource :: proc(state: ^Resource_State, b: ^bk.Backend, id: u32) {
	if id == 0 || id >= MAX_COMPUTE_SHADERS || !state.compute_shaders[id].active {
		return
	}
	cs := &state.compute_shaders[id]
	b.destroy_descriptor_pool(cs.desc_pool)
	b.destroy_descriptor_set_layout(cs.desc_layout)
	b.destroy_compute_pipeline(cs.pipeline)
	b.destroy_shader(cs.shader)
	cs^ = {}
}

// --- Storage Buffer Resources ---

create_storage_buffer_resource :: proc(
	state: ^Resource_State,
	b: ^bk.Backend,
	size: int,
	data: rawptr = nil,
) -> (id: u32, ok: bool) {
	// Allocate slot
	slot: u32
	slot_found := false
	for i in 1..<u32(MAX_STORAGE_BUFFERS) {
		if !state.storage_buffers[i].active {
			slot = i
			slot_found = true
			break
		}
	}
	if !slot_found {
		log.error("gpu/resource: storage buffer pool exhausted")
		return 0, false
	}

	// Create host-visible storage buffer
	buf, buf_ok := b.create_buffer(bk.Buffer_Desc{
		size   = u64(size),
		usage  = {.Storage},
		memory = {.Host_Visible, .Host_Coherent},
	})
	if !buf_ok {
		return 0, false
	}

	// Persistently map
	mapped := b.map_buffer(buf)

	// Copy initial data if provided
	if data != nil && mapped != nil {
		mem.copy(mapped, data, size)
	}

	state.storage_buffers[slot] = Internal_Storage_Buffer{
		buffer = buf,
		size   = size,
		active = true,
	}

	return slot, true
}

unload_storage_buffer_resource :: proc(state: ^Resource_State, b: ^bk.Backend, id: u32) {
	if id == 0 || id >= MAX_STORAGE_BUFFERS || !state.storage_buffers[id].active {
		return
	}
	sb := &state.storage_buffers[id]
	b.unmap_buffer(sb.buffer)
	b.destroy_buffer(sb.buffer)
	sb.active = false
}

get_storage_buffer :: proc(state: ^Resource_State, id: u32) -> (^Internal_Storage_Buffer, bool) {
	if id == 0 || id >= MAX_STORAGE_BUFFERS || !state.storage_buffers[id].active {
		return nil, false
	}
	return &state.storage_buffers[id], true
}

get_compute_shader :: proc(state: ^Resource_State, id: u32) -> (^Internal_Compute_Shader, bool) {
	if id == 0 || id >= MAX_COMPUTE_SHADERS || !state.compute_shaders[id].active {
		return nil, false
	}
	return &state.compute_shaders[id], true
}

shutdown_compute :: proc(state: ^Resource_State, b: ^bk.Backend) {
	for i in 0..<u32(MAX_COMPUTE_SHADERS) {
		if state.compute_shaders[i].active {
			unload_compute_shader_resource(state, b, i)
		}
	}
	for i in 0..<u32(MAX_STORAGE_BUFFERS) {
		if state.storage_buffers[i].active {
			unload_storage_buffer_resource(state, b, i)
		}
	}
}

@(private)
destroy_texture_slot :: proc(b: ^bk.Backend, tex: ^Internal_Texture) {
	// Descriptor set freed when pool is destroyed
	if tex.sampler != bk.NULL_SAMPLER {
		b.destroy_sampler(tex.sampler)
		tex.sampler = bk.NULL_SAMPLER
	}
	if tex.texture != bk.NULL_TEXTURE {
		if tex.owns_texture {
			b.destroy_texture(tex.texture)
		}
		tex.texture = bk.NULL_TEXTURE
	}
	tex.active = false
}