//! Vulkan device enumeration with a discrete-compute filter.
//!
//! The goal is to produce a `BoxReportPayload` skeleton for one host:
//! the list of bus addresses present, per-device driver versions, and
//! the host identity. Real measurements come later. The filter rejects
//! anything that isn't a discrete GPU with a non-trivial device-local
//! heap, which keeps BMC display chips (ASPEED, Matrox) and integrated
//! GPUs out of the list without ever consulting a vendor table.

use std::ffi::CStr;

use ash::{Entry, Instance, vk};
use everycycle_hal::{BoxReportPayload, DriverBinding, HostIdentity};

use crate::host::host_identity;

/// One discrete compute device the enumeration found.
#[derive(Clone, Debug)]
pub struct EnumeratedDevice {
    pub bus_address: String,
    pub pci_ids: String,
    pub driver_version: String,
    pub device_name: String,
    pub physical_bytes: u64,
}

/// Result of enumerating one host.
#[derive(Clone, Debug)]
pub struct EnumeratedHost {
    pub host: HostIdentity,
    pub devices: Vec<EnumeratedDevice>,
}

impl EnumeratedHost {
    /// Build a `BoxReportPayload` skeleton from this enumeration.
    ///
    /// `probed_at` should be an RFC 3339 timestamp; the caller supplies
    /// it because there are several sensible time sources and no
    /// reason for this layer to choose. `ambient_c` is similarly the
    /// caller's responsibility once a thermal probe lands.
    #[must_use]
    pub fn into_box_payload(
        self,
        probed_at: impl Into<String>,
        ambient_c: f32,
        probe_suite_version: impl Into<String>,
    ) -> BoxReportPayload {
        let present_devices = self.devices.iter().map(|d| d.bus_address.clone()).collect();
        let driver_versions = self
            .devices
            .iter()
            .map(|d| DriverBinding {
                bus_address: d.bus_address.clone(),
                driver_version: d.driver_version.clone(),
            })
            .collect();
        BoxReportPayload {
            report_version: "1".into(),
            probed_at: probed_at.into(),
            host: self.host,
            present_devices,
            driver_versions,
            ambient_c,
            probe_suite_version: probe_suite_version.into(),
        }
    }
}

#[derive(Debug)]
pub enum EnumError {
    VulkanLoad(ash::LoadingError),
    Vulkan(vk::Result),
}

impl core::fmt::Display for EnumError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::VulkanLoad(e) => write!(f, "could not load Vulkan loader: {e:?}"),
            Self::Vulkan(r) => write!(f, "Vulkan call failed: {r:?}"),
        }
    }
}

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

const MIN_DEVICE_LOCAL_BYTES: u64 = 1_000_000_000;

/// Enumerate all discrete compute devices on the current host.
pub fn enumerate() -> Result<EnumeratedHost, EnumError> {
    let entry = unsafe { Entry::load() }.map_err(EnumError::VulkanLoad)?;
    let app_info = vk::ApplicationInfo::default().api_version(vk::API_VERSION_1_3);
    let create_info = vk::InstanceCreateInfo::default().application_info(&app_info);
    let instance =
        unsafe { entry.create_instance(&create_info, None) }.map_err(EnumError::Vulkan)?;

    let result = (|| -> Result<EnumeratedHost, EnumError> {
        let devices =
            unsafe { instance.enumerate_physical_devices() }.map_err(EnumError::Vulkan)?;
        let mut out = Vec::new();
        for dev in devices {
            if let Some(d) = inspect_device(&instance, dev) {
                out.push(d);
            }
        }
        Ok(EnumeratedHost {
            host: host_identity(),
            devices: out,
        })
    })();

    unsafe { instance.destroy_instance(None) };
    result
}

fn inspect_device(instance: &Instance, dev: vk::PhysicalDevice) -> Option<EnumeratedDevice> {
    let props = unsafe { instance.get_physical_device_properties(dev) };
    if props.device_type != vk::PhysicalDeviceType::DISCRETE_GPU {
        return None;
    }

    let mem = unsafe { instance.get_physical_device_memory_properties(dev) };
    let physical_bytes = mem
        .memory_heaps
        .iter()
        .take(mem.memory_heap_count as usize)
        .filter(|h| h.flags.contains(vk::MemoryHeapFlags::DEVICE_LOCAL))
        .map(|h| h.size)
        .max()
        .unwrap_or(0);
    if physical_bytes < MIN_DEVICE_LOCAL_BYTES {
        return None;
    }

    let device_name = unsafe { CStr::from_ptr(props.device_name.as_ptr()) }
        .to_string_lossy()
        .into_owned();
    // PCI vendor and device IDs are 16-bit by spec; Vulkan reports them
    // widened to u32. Format as 4 hex digits without truncating casts.
    let pci_ids = format!("{:04x}:{:04x}", props.vendor_id, props.device_id);

    let bus_address = pci_bus_address(instance, dev).unwrap_or_else(|| "0000:00:00.0".to_string());
    let driver_version = driver_string(instance, dev);

    Some(EnumeratedDevice {
        bus_address,
        pci_ids,
        driver_version,
        device_name,
        physical_bytes,
    })
}

fn pci_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 driver_string(instance: &Instance, dev: vk::PhysicalDevice) -> String {
    let mut driver = vk::PhysicalDeviceDriverProperties::default();
    let mut props2 = vk::PhysicalDeviceProperties2::default().push_next(&mut driver);
    unsafe { instance.get_physical_device_properties2(dev, &mut props2) };
    let name = unsafe { CStr::from_ptr(driver.driver_name.as_ptr()) }
        .to_string_lossy()
        .into_owned();
    let info = unsafe { CStr::from_ptr(driver.driver_info.as_ptr()) }
        .to_string_lossy()
        .into_owned();
    match (name.is_empty(), info.is_empty()) {
        (true, true) => "unknown".to_string(),
        (false, true) => name,
        (true, false) => info,
        (false, false) => format!("{name} ({info})"),
    }
}

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 name = unsafe { CStr::from_ptr(ext.extension_name.as_ptr()) };
        name == want
    })
}