//! Shared Vulkan setup for probes: instance, physical device,
//! logical device, compute queue, command pool.
//!
//! A `VulkanContext` is created once per appraisal run, against one
//! specific physical device identified by its PCI bus address. Probes
//! borrow it to record command buffers and submit work; resource
//! cleanup happens when the context is dropped.

use std::ffi::{CStr, c_char};

use ash::{Entry, Instance, vk};

const APP_NAME: &CStr = c"everycycle-appraise";

#[derive(Debug)]
pub enum VulkanContextError {
    Load(ash::LoadingError),
    Vulkan(vk::Result),
    NoMatchingDevice(String),
    NoComputeQueueFamily,
}

impl core::fmt::Display for VulkanContextError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::Load(e) => write!(f, "could not load Vulkan loader: {e:?}"),
            Self::Vulkan(r) => write!(f, "Vulkan call failed: {r:?}"),
            Self::NoMatchingDevice(addr) => {
                write!(f, "no physical device at bus address {addr}")
            }
            Self::NoComputeQueueFamily => {
                write!(f, "no queue family on the selected device supports COMPUTE")
            }
        }
    }
}

impl std::error::Error for VulkanContextError {}

pub struct VulkanContext {
    pub queue_family_index: u32,
    pub queue: vk::Queue,
    pub command_pool: vk::CommandPool,
    pub physical_device: vk::PhysicalDevice,
    pub memory_properties: vk::PhysicalDeviceMemoryProperties,

    /// Nanoseconds per timestamp-query tick on this device.
    pub timestamp_period_ns: f32,
    /// Meaningful low-order bits in a timestamp query on the chosen
    /// queue family. Zero means the queue cannot write timestamps, and
    /// probes must fall back to host timing.
    pub timestamp_valid_bits: u32,

    // Held for cleanup. Dropped in reverse declaration order.
    pub device: ash::Device,
    pub instance: Instance,
    _entry: Entry,
}

impl VulkanContext {
    /// Pick the physical device whose PCI bus address matches `target`
    /// and bring up the logical device, compute queue, and command pool
    /// against it.
    pub fn new(target_bus_address: &str) -> Result<Self, VulkanContextError> {
        let entry = unsafe { Entry::load() }.map_err(VulkanContextError::Load)?;
        let instance = create_instance(&entry)?;

        let physical = pick_device(&instance, target_bus_address)?;
        let memory_properties = unsafe { instance.get_physical_device_memory_properties(physical) };

        let queue_family_index = pick_compute_queue_family(&instance, physical)?;

        let timestamp_period_ns = unsafe { instance.get_physical_device_properties(physical) }
            .limits
            .timestamp_period;
        let timestamp_valid_bits =
            unsafe { instance.get_physical_device_queue_family_properties(physical) }
                .get(queue_family_index as usize)
                .map_or(0, |fam| fam.timestamp_valid_bits);

        let priorities = [1.0_f32];
        let queue_info = vk::DeviceQueueCreateInfo::default()
            .queue_family_index(queue_family_index)
            .queue_priorities(&priorities);
        let device_info =
            vk::DeviceCreateInfo::default().queue_create_infos(std::slice::from_ref(&queue_info));
        let device = unsafe { instance.create_device(physical, &device_info, None) }
            .map_err(VulkanContextError::Vulkan)?;

        let queue = unsafe { device.get_device_queue(queue_family_index, 0) };

        let pool_info = vk::CommandPoolCreateInfo::default()
            .flags(vk::CommandPoolCreateFlags::RESET_COMMAND_BUFFER)
            .queue_family_index(queue_family_index);
        let command_pool = unsafe { device.create_command_pool(&pool_info, None) }
            .map_err(VulkanContextError::Vulkan)?;

        Ok(Self {
            queue_family_index,
            queue,
            command_pool,
            physical_device: physical,
            memory_properties,
            timestamp_period_ns,
            timestamp_valid_bits,
            device,
            instance,
            _entry: entry,
        })
    }

    /// Find a memory type index matching `required_flags` from the
    /// device's memory-properties table and a bitmask of acceptable
    /// types (`type_filter`, as returned by `MemoryRequirements`).
    #[must_use]
    pub fn find_memory_type(
        &self,
        type_filter: u32,
        required_flags: vk::MemoryPropertyFlags,
    ) -> Option<u32> {
        for i in 0..self.memory_properties.memory_type_count {
            let supported = (type_filter & (1 << i)) != 0;
            let mt = self.memory_properties.memory_types[i as usize];
            if supported && mt.property_flags.contains(required_flags) {
                return Some(i);
            }
        }
        None
    }
}

impl Drop for VulkanContext {
    fn drop(&mut self) {
        unsafe {
            // Wait for any outstanding work before tearing things down.
            // Errors at drop time can't be returned; the destructor
            // calls themselves are infallible.
            let _ = self.device.device_wait_idle();
            self.device.destroy_command_pool(self.command_pool, None);
            self.device.destroy_device(None);
            self.instance.destroy_instance(None);
        }
    }
}

fn create_instance(entry: &Entry) -> Result<Instance, VulkanContextError> {
    let app_info = vk::ApplicationInfo::default()
        .application_name(APP_NAME)
        .engine_name(APP_NAME)
        .api_version(vk::API_VERSION_1_3);
    let create_info = vk::InstanceCreateInfo::default().application_info(&app_info);
    unsafe { entry.create_instance(&create_info, None) }.map_err(VulkanContextError::Vulkan)
}

fn pick_device(
    instance: &Instance,
    target: &str,
) -> Result<vk::PhysicalDevice, VulkanContextError> {
    let devices =
        unsafe { instance.enumerate_physical_devices() }.map_err(VulkanContextError::Vulkan)?;
    for dev in devices {
        if device_bus_address(instance, dev).as_deref() == Some(target) {
            return Ok(dev);
        }
    }
    Err(VulkanContextError::NoMatchingDevice(target.to_string()))
}

fn device_bus_address(instance: &Instance, dev: vk::PhysicalDevice) -> Option<String> {
    if !device_supports_extension(instance, dev, ash::ext::pci_bus_info::NAME) {
        return None;
    }
    let mut pci = vk::PhysicalDevicePCIBusInfoPropertiesEXT::default();
    let mut props2 = vk::PhysicalDeviceProperties2::default().push_next(&mut pci);
    unsafe { instance.get_physical_device_properties2(dev, &mut props2) };
    Some(format!(
        "{:04x}:{:02x}:{:02x}.{:x}",
        pci.pci_domain, pci.pci_bus, pci.pci_device, pci.pci_function
    ))
}

fn device_supports_extension(instance: &Instance, dev: vk::PhysicalDevice, want: &CStr) -> bool {
    let Ok(exts) = (unsafe { instance.enumerate_device_extension_properties(dev) }) else {
        return false;
    };
    exts.iter().any(|ext| {
        let raw = ext.extension_name.as_ptr().cast::<c_char>();
        let name = unsafe { CStr::from_ptr(raw) };
        name == want
    })
}

fn pick_compute_queue_family(
    instance: &Instance,
    physical: vk::PhysicalDevice,
) -> Result<u32, VulkanContextError> {
    let families = unsafe { instance.get_physical_device_queue_family_properties(physical) };
    for (i, fam) in families.iter().enumerate() {
        if fam.queue_flags.contains(vk::QueueFlags::COMPUTE) {
            #[allow(clippy::cast_possible_truncation)]
            return Ok(i as u32);
        }
    }
    Err(VulkanContextError::NoComputeQueueFamily)
}