package gpu_core

import "core:log"
import vk "vendor:vulkan"

Queue_Families :: struct {
	graphics: u32,
	present:  u32,
	compute:  u32,
	has_graphics: bool,
	has_present:  bool,
	has_compute:  bool,
}

Gpu_Device :: struct {
	physical_device: vk.PhysicalDevice,
	device:          vk.Device,
	queue_families:  Queue_Families,
	graphics_queue:  vk.Queue,
	present_queue:   vk.Queue,
	compute_queue:   vk.Queue,
	properties:      vk.PhysicalDeviceProperties,
}

@(private)
required_device_extensions := [?]cstring{
	vk.KHR_SWAPCHAIN_EXTENSION_NAME,
}

create_device :: proc(instance: vk.Instance, surface: vk.SurfaceKHR) -> (dev: Gpu_Device, ok: bool) {
	// Enumerate physical devices
	device_count: u32
	vk.EnumeratePhysicalDevices(instance, &device_count, nil)
	if device_count == 0 {
		log.error("gpu/core: no Vulkan-capable GPU found")
		return {}, false
	}

	devices := make([]vk.PhysicalDevice, device_count, context.temp_allocator)
	vk.EnumeratePhysicalDevices(instance, &device_count, raw_data(devices))

	// Pick best device (prefer discrete GPU)
	best_score := -1
	for d in devices {
		score := rate_device(d, surface)
		if score > best_score {
			best_score = score
			dev.physical_device = d
		}
	}

	if best_score < 0 {
		log.error("gpu/core: no suitable GPU found")
		return {}, false
	}

	vk.GetPhysicalDeviceProperties(dev.physical_device, &dev.properties)
	log.infof("gpu/core: selected GPU: %s", dev.properties.deviceName)

	// Find queue families
	dev.queue_families = find_queue_families(dev.physical_device, surface)
	if !dev.queue_families.has_graphics || !dev.queue_families.has_present {
		log.error("gpu/core: GPU missing required queue families")
		return {}, false
	}

	// Create logical device
	unique_families: [3]u32
	unique_count: u32 = 0

	add_unique :: proc(arr: ^[3]u32, count: ^u32, val: u32) {
		for i in 0..<count^ {
			if arr[i] == val {
				return
			}
		}
		arr[count^] = val
		count^ += 1
	}

	add_unique(&unique_families, &unique_count, dev.queue_families.graphics)
	add_unique(&unique_families, &unique_count, dev.queue_families.present)
	if dev.queue_families.has_compute {
		add_unique(&unique_families, &unique_count, dev.queue_families.compute)
	}

	queue_priority: f32 = 1.0
	queue_create_infos: [3]vk.DeviceQueueCreateInfo
	for i in 0..<unique_count {
		queue_create_infos[i] = vk.DeviceQueueCreateInfo{
			sType            = .DEVICE_QUEUE_CREATE_INFO,
			queueFamilyIndex = unique_families[i],
			queueCount       = 1,
			pQueuePriorities = &queue_priority,
		}
	}

	features := vk.PhysicalDeviceFeatures{
		samplerAnisotropy = true,
		fillModeNonSolid  = true,
	}

	device_info := vk.DeviceCreateInfo{
		sType                   = .DEVICE_CREATE_INFO,
		queueCreateInfoCount    = unique_count,
		pQueueCreateInfos       = &queue_create_infos[0],
		enabledExtensionCount   = u32(len(required_device_extensions)),
		ppEnabledExtensionNames = &required_device_extensions[0],
		pEnabledFeatures        = &features,
	}

	result := vk.CreateDevice(dev.physical_device, &device_info, nil, &dev.device)
	if result != .SUCCESS {
		log.errorf("gpu/core: vkCreateDevice failed: %v", result)
		return {}, false
	}

	// Load device-level procedures
	vk.load_proc_addresses_device(dev.device)

	// Get queues
	vk.GetDeviceQueue(dev.device, dev.queue_families.graphics, 0, &dev.graphics_queue)
	vk.GetDeviceQueue(dev.device, dev.queue_families.present, 0, &dev.present_queue)
	if dev.queue_families.has_compute {
		vk.GetDeviceQueue(dev.device, dev.queue_families.compute, 0, &dev.compute_queue)
	}

	return dev, true
}

create_device_headless :: proc(instance: vk.Instance) -> (dev: Gpu_Device, ok: bool) {
	device_count: u32
	vk.EnumeratePhysicalDevices(instance, &device_count, nil)
	if device_count == 0 {
		log.error("gpu/core: no Vulkan-capable GPU found")
		return {}, false
	}

	devices := make([]vk.PhysicalDevice, device_count, context.temp_allocator)
	vk.EnumeratePhysicalDevices(instance, &device_count, raw_data(devices))

	best_score := -1
	for d in devices {
		score := rate_device_headless(d)
		if score > best_score {
			best_score = score
			dev.physical_device = d
		}
	}

	if best_score < 0 {
		log.error("gpu/core: no suitable headless GPU found")
		return {}, false
	}

	vk.GetPhysicalDeviceProperties(dev.physical_device, &dev.properties)
	log.infof("gpu/core: selected headless GPU: %s", dev.properties.deviceName)

	dev.queue_families = find_queue_families_headless(dev.physical_device)
	if !dev.queue_families.has_graphics || !dev.queue_families.has_compute {
		log.error("gpu/core: headless GPU missing graphics or compute queue")
		return {}, false
	}

	unique_families: [2]u32
	unique_count: u32 = 0

	add_unique :: proc(arr: ^[2]u32, count: ^u32, val: u32) {
		for i in 0..<count^ {
			if arr[i] == val do return
		}
		arr[count^] = val
		count^ += 1
	}

	add_unique(&unique_families, &unique_count, dev.queue_families.graphics)
	add_unique(&unique_families, &unique_count, dev.queue_families.compute)

	queue_priority: f32 = 1.0
	queue_create_infos: [2]vk.DeviceQueueCreateInfo
	for i in 0..<unique_count {
		queue_create_infos[i] = vk.DeviceQueueCreateInfo{
			sType            = .DEVICE_QUEUE_CREATE_INFO,
			queueFamilyIndex = unique_families[i],
			queueCount       = 1,
			pQueuePriorities = &queue_priority,
		}
	}

	features := vk.PhysicalDeviceFeatures{
		samplerAnisotropy = true,
		fillModeNonSolid  = true,
	}

	device_info := vk.DeviceCreateInfo{
		sType                   = .DEVICE_CREATE_INFO,
		queueCreateInfoCount    = unique_count,
		pQueueCreateInfos       = &queue_create_infos[0],
		enabledExtensionCount   = u32(len(required_device_extensions)),
		ppEnabledExtensionNames = &required_device_extensions[0],
		pEnabledFeatures        = &features,
	}

	result := vk.CreateDevice(dev.physical_device, &device_info, nil, &dev.device)
	if result != .SUCCESS {
		log.errorf("gpu/core: headless vkCreateDevice failed: %v", result)
		return {}, false
	}

	vk.load_proc_addresses_device(dev.device)
	vk.GetDeviceQueue(dev.device, dev.queue_families.graphics, 0, &dev.graphics_queue)
	vk.GetDeviceQueue(dev.device, dev.queue_families.compute, 0, &dev.compute_queue)
	dev.present_queue = nil
	dev.queue_families.present = dev.queue_families.graphics
	dev.queue_families.has_present = false

	return dev, true
}

destroy_device :: proc(dev: ^Gpu_Device) {
	if dev.device != nil {
		vk.DestroyDevice(dev.device, nil)
		dev.device = nil
	}
}

@(private)
find_queue_families :: proc(device: vk.PhysicalDevice, surface: vk.SurfaceKHR) -> Queue_Families {
	families: Queue_Families

	count: u32
	vk.GetPhysicalDeviceQueueFamilyProperties(device, &count, nil)
	props := make([]vk.QueueFamilyProperties, count, context.temp_allocator)
	vk.GetPhysicalDeviceQueueFamilyProperties(device, &count, raw_data(props))

	for p, i in props {
		idx := u32(i)

		if .GRAPHICS in p.queueFlags && !families.has_graphics {
			families.graphics = idx
			families.has_graphics = true
		}

		if .COMPUTE in p.queueFlags && !families.has_compute {
			families.compute = idx
			families.has_compute = true
		}

		present_support: b32
		vk.GetPhysicalDeviceSurfaceSupportKHR(device, idx, surface, &present_support)
		if present_support && !families.has_present {
			families.present = idx
			families.has_present = true
		}
	}

	return families
}

@(private)
find_queue_families_headless :: proc(device: vk.PhysicalDevice) -> Queue_Families {
	families: Queue_Families

	count: u32
	vk.GetPhysicalDeviceQueueFamilyProperties(device, &count, nil)
	props := make([]vk.QueueFamilyProperties, count, context.temp_allocator)
	vk.GetPhysicalDeviceQueueFamilyProperties(device, &count, raw_data(props))

	for p, i in props {
		idx := u32(i)
		if .GRAPHICS in p.queueFlags && !families.has_graphics {
			families.graphics = idx
			families.has_graphics = true
		}
		if .COMPUTE in p.queueFlags && !families.has_compute {
			families.compute = idx
			families.has_compute = true
		}
	}

	return families
}

@(private)
rate_device :: proc(device: vk.PhysicalDevice, surface: vk.SurfaceKHR) -> int {
	props: vk.PhysicalDeviceProperties
	vk.GetPhysicalDeviceProperties(device, &props)

	// Check required extensions
	ext_count: u32
	vk.EnumerateDeviceExtensionProperties(device, nil, &ext_count, nil)
	exts := make([]vk.ExtensionProperties, ext_count, context.temp_allocator)
	vk.EnumerateDeviceExtensionProperties(device, nil, &ext_count, raw_data(exts))

	for required in required_device_extensions {
		found := false
		for &ext in exts {
			name := cstring(raw_data(&ext.extensionName))
			if name == required {
				found = true
				break
			}
		}
		if !found {
			return -1
		}
	}

	// Check queue families
	families := find_queue_families(device, surface)
	if !families.has_graphics || !families.has_present {
		return -1
	}

	// Check swapchain support
	format_count: u32
	vk.GetPhysicalDeviceSurfaceFormatsKHR(device, surface, &format_count, nil)
	present_count: u32
	vk.GetPhysicalDeviceSurfacePresentModesKHR(device, surface, &present_count, nil)
	if format_count == 0 || present_count == 0 {
		return -1
	}

	// Score
	score := 0
	if props.deviceType == .DISCRETE_GPU {
		score += 1000
	} else if props.deviceType == .INTEGRATED_GPU {
		score += 100
	}

	return score
}

@(private)
rate_device_headless :: proc(device: vk.PhysicalDevice) -> int {
	props: vk.PhysicalDeviceProperties
	vk.GetPhysicalDeviceProperties(device, &props)

	families := find_queue_families_headless(device)
	if !families.has_graphics || !families.has_compute {
		return -1
	}
	ext_count: u32
	vk.EnumerateDeviceExtensionProperties(device, nil, &ext_count, nil)
	exts := make([]vk.ExtensionProperties, ext_count, context.temp_allocator)
	vk.EnumerateDeviceExtensionProperties(device, nil, &ext_count, raw_data(exts))
	for required in required_device_extensions {
		found := false
		for &ext in exts {
			name := cstring(raw_data(&ext.extensionName))
			if name == required {
				found = true
				break
			}
		}
		if !found {
			return -1
		}
	}

	score := 0
	if props.deviceType == .DISCRETE_GPU {
		score += 1000
	} else if props.deviceType == .INTEGRATED_GPU {
		score += 100
	}
	return score
}