singe-cupti 0.1.0-alpha.5

use std::{
    ffi::{CStr, CString, c_char},
    marker::PhantomData,
    mem::{MaybeUninit, size_of, zeroed},
    panic::{AssertUnwindSafe, catch_unwind},
    ptr,
    sync::Mutex,
};

use singe_cuda::device::Device;
use singe_cupti_sys as sys;

use crate::{
    context::Context,
    error::{Error, Result},
    try_ffi,
    types::{
        DeviceAttribute, DeviceAttributeDeviceClass, DeviceAttributeValue, EventAttribute,
        EventAttributeValue, EventCategory, EventCollectionMethod, EventCollectionMode,
        EventDomainAttribute, EventDomainAttributeValue, EventDomainId, EventGroupAttribute,
        EventGroupAttributeSetting, EventGroupAttributeValue, EventId, EventProfilingScope,
        MetricAttribute, MetricAttributeValue, MetricCategory, MetricEvaluationMode, MetricId,
        MetricPropertyId, MetricValue, MetricValueKind, ReadEventFlags,
    },
    utility::{checked_byte_len, to_u64, to_usize},
};

type KernelReplayUpdateCallback = dyn FnMut(&CStr, i32) + Send + 'static;

static KERNEL_REPLAY_UPDATE_CALLBACK: Mutex<Option<Box<KernelReplayUpdateCallback>>> =
    Mutex::new(None);

fn query_array<T: Copy>(
    mut query: impl FnMut(*mut u64, *mut T) -> sys::CUptiResult,
) -> Result<Vec<T>> {
    let mut size_bytes = 0;
    try_ffi!(query(&mut size_bytes, ptr::null_mut()))?;

    if size_bytes == 0 {
        return Ok(Vec::new());
    }

    let element_size = size_of::<T>();
    let size_bytes = to_usize(size_bytes, "array_size_bytes")?;

    if size_bytes % element_size != 0 {
        return Err(Error::LengthMismatch {
            name: "array_size_bytes".to_string(),
            expected: element_size,
            actual: size_bytes,
        });
    }

    let len = size_bytes / element_size;
    let mut values = vec![unsafe { zeroed() }; len];
    let mut retry_size_bytes = size_bytes as u64;
    try_ffi!(query(&mut retry_size_bytes, values.as_mut_ptr()))?;
    let retry_size_bytes = to_usize(retry_size_bytes, "array_size_bytes")?;
    values.truncate(retry_size_bytes / element_size);
    Ok(values)
}

#[derive(Debug)]
pub struct EventGroup {
    handle: sys::CUpti_EventGroup,
}

#[derive(Debug)]
pub struct EventGroupSets {
    handle: *mut sys::CUpti_EventGroupSets,
}

#[derive(Debug, Clone, Copy)]
pub struct EventGroupSet<'a> {
    handle: *mut sys::CUpti_EventGroupSet,
    _marker: PhantomData<&'a EventGroupSets>,
}

#[derive(Debug, Clone)]
pub struct EventGroupReadAll {
    pub event_ids: Vec<EventId>,
    pub values: Vec<u64>,
}

impl EventGroup {
    fn from_raw(handle: sys::CUpti_EventGroup) -> Result<Self> {
        if handle.is_null() {
            return Err(Error::NullHandle);
        }
        Ok(Self { handle })
    }

    pub(crate) const fn as_raw(&self) -> sys::CUpti_EventGroup {
        self.handle
    }
}

impl Drop for EventGroup {
    fn drop(&mut self) {
        unsafe {
            let _ = sys::cuptiEventGroupDestroy(self.handle);
        }
    }
}

impl EventGroupSets {
    fn from_raw(handle: *mut sys::CUpti_EventGroupSets) -> Result<Self> {
        if handle.is_null() {
            return Err(Error::NullHandle);
        }
        Ok(Self { handle })
    }

    pub fn sets(&self) -> Vec<EventGroupSet<'_>> {
        let sets = unsafe { &*self.handle };
        if sets.sets.is_null() || sets.numSets == 0 {
            return Vec::new();
        }
        (0..sets.numSets as usize)
            .map(|index| EventGroupSet {
                handle: unsafe { sets.sets.add(index) },
                _marker: PhantomData,
            })
            .collect()
    }
}

impl Drop for EventGroupSets {
    fn drop(&mut self) {
        unsafe {
            let _ = sys::cuptiEventGroupSetsDestroy(self.handle);
        }
    }
}

impl EventGroupSet<'_> {
    pub(crate) const fn as_raw(self) -> *mut sys::CUpti_EventGroupSet {
        self.handle
    }

    pub fn enable(self) -> Result<()> {
        unsafe {
            try_ffi!(sys::cuptiEventGroupSetEnable(self.as_raw()))?;
        }
        Ok(())
    }

    pub fn disable(self) -> Result<()> {
        unsafe {
            try_ffi!(sys::cuptiEventGroupSetDisable(self.as_raw()))?;
        }
        Ok(())
    }
}

pub fn set_event_collection_mode(context: &Context, mode: EventCollectionMode) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiSetEventCollectionMode(
            context.as_raw(),
            mode.into()
        ))?;
    }
    Ok(())
}

fn raw_device_get_attribute(
    device: sys::CUdevice,
    attrib: DeviceAttribute,
    value_size: *mut u64,
    value: *mut (),
) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiDeviceGetAttribute(
            device,
            attrib.into(),
            value_size,
            value as _,
        ))?;
    }
    Ok(())
}

fn raw_device_num_event_domains(device: sys::CUdevice, num_domains: *mut u32) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiDeviceGetNumEventDomains(device, num_domains))?;
    }
    Ok(())
}

fn read_attribute_scalar<T: Copy>(
    mut read: impl FnMut(*mut u64, *mut ()) -> Result<()>,
    name: &'static str,
) -> Result<T> {
    let mut value = MaybeUninit::<T>::uninit();
    let mut value_size = to_u64(size_of::<T>(), name)?;
    read(&mut value_size, value.as_mut_ptr().cast())?;
    Ok(unsafe { value.assume_init() })
}

fn read_attribute_string(
    mut read: impl FnMut(*mut u64, *mut ()) -> Result<()>,
    name: &'static str,
) -> Result<String> {
    let mut buffer = vec![0u8; 4096];
    let mut value_size = to_u64(buffer.len(), name)?;
    read(&mut value_size, buffer.as_mut_ptr().cast())?;
    let len = buffer
        .iter()
        .position(|byte| *byte == 0)
        .unwrap_or_else(|| value_size.min(buffer.len() as u64) as usize);
    Ok(String::from_utf8_lossy(&buffer[..len]).into_owned())
}

pub fn device_attribute_value(
    device: Device,
    attrib: DeviceAttribute,
) -> Result<DeviceAttributeValue> {
    let mut read = |value_size: *mut u64, value: *mut ()| {
        raw_device_get_attribute(device.id() as sys::CUdevice, attrib, value_size, value)
    };
    match attrib {
        DeviceAttribute::DeviceClass => Ok(DeviceAttributeValue::DeviceClass(
            DeviceAttributeDeviceClass::from(read_attribute_scalar::<
                sys::CUpti_DeviceAttributeDeviceClass,
            >(&mut read, "device_attribute_size")?),
        )),
        DeviceAttribute::ForceInt => Err(Error::InvalidAttribute {
            name: format!("{attrib:?}"),
        }),
        _ => Ok(DeviceAttributeValue::U32(read_attribute_scalar::<u32>(
            &mut read,
            "device_attribute_size",
        )?)),
    }
}

pub fn device_num_event_domains(device: Device) -> Result<u32> {
    let mut num_domains = 0;
    raw_device_num_event_domains(device.id() as sys::CUdevice, &mut num_domains)?;
    Ok(num_domains)
}

pub fn device_enum_event_domains(device: Device) -> Result<Vec<EventDomainId>> {
    let raw = query_array(|array_size_bytes, domain_array| unsafe {
        sys::cuptiDeviceEnumEventDomains(
            device.id() as sys::CUdevice,
            array_size_bytes,
            domain_array,
        )
    })?;
    Ok(raw.into_iter().map(EventDomainId::from).collect())
}

fn raw_device_get_event_domain_attribute(
    device: sys::CUdevice,
    event_domain: EventDomainId,
    attrib: EventDomainAttribute,
    value_size: *mut u64,
    value: *mut (),
) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiDeviceGetEventDomainAttribute(
            device,
            event_domain.as_raw(),
            attrib.into(),
            value_size,
            value as _,
        ))?;
    }
    Ok(())
}

pub fn device_event_domain_attribute_value(
    device: Device,
    event_domain: EventDomainId,
    attrib: EventDomainAttribute,
) -> Result<EventDomainAttributeValue> {
    let mut read = |value_size: *mut u64, value: *mut ()| {
        raw_device_get_event_domain_attribute(
            device.id() as sys::CUdevice,
            event_domain,
            attrib,
            value_size,
            value,
        )
    };
    event_domain_attribute_value_with(&mut read, attrib)
}

fn raw_num_event_domains(num_domains: *mut u32) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiGetNumEventDomains(num_domains))?;
    }
    Ok(())
}

pub fn num_event_domains() -> Result<u32> {
    let mut num_domains = 0;
    raw_num_event_domains(&mut num_domains)?;
    Ok(num_domains)
}

pub fn enum_event_domains() -> Result<Vec<EventDomainId>> {
    let raw = query_array(|array_size_bytes, domain_array| unsafe {
        sys::cuptiEnumEventDomains(array_size_bytes, domain_array)
    })?;
    Ok(raw.into_iter().map(EventDomainId::from).collect())
}

fn raw_event_domain_get_attribute(
    event_domain: sys::CUpti_EventDomainID,
    attrib: EventDomainAttribute,
    value_size: *mut u64,
    value: *mut (),
) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiEventDomainGetAttribute(
            event_domain,
            attrib.into(),
            value_size,
            value as _,
        ))?;
    }
    Ok(())
}

fn raw_event_domain_num_events(
    event_domain: sys::CUpti_EventDomainID,
    num_events: *mut u32,
) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiEventDomainGetNumEvents(event_domain, num_events))?;
    }
    Ok(())
}

pub fn event_domain_attribute_value(
    event_domain: EventDomainId,
    attrib: EventDomainAttribute,
) -> Result<EventDomainAttributeValue> {
    let mut read = |value_size: *mut u64, value: *mut ()| {
        raw_event_domain_get_attribute(event_domain.as_raw(), attrib, value_size, value)
    };
    event_domain_attribute_value_with(&mut read, attrib)
}

fn event_domain_attribute_value_with(
    read: &mut impl FnMut(*mut u64, *mut ()) -> Result<()>,
    attrib: EventDomainAttribute,
) -> Result<EventDomainAttributeValue> {
    match attrib {
        EventDomainAttribute::Name => Ok(EventDomainAttributeValue::Name(read_attribute_string(
            read,
            "event_domain_attribute_size",
        )?)),
        EventDomainAttribute::InstanceCount | EventDomainAttribute::TotalInstanceCount => {
            Ok(EventDomainAttributeValue::U32(
                read_attribute_scalar::<u32>(read, "event_domain_attribute_size")?,
            ))
        }
        EventDomainAttribute::CollectionMethod => Ok(EventDomainAttributeValue::CollectionMethod(
            EventCollectionMethod::from(read_attribute_scalar::<sys::CUpti_EventCollectionMethod>(
                read,
                "event_domain_attribute_size",
            )?),
        )),
        EventDomainAttribute::ForceInt => Err(Error::InvalidAttribute {
            name: format!("{attrib:?}"),
        }),
    }
}

pub fn event_domain_num_events(event_domain: EventDomainId) -> Result<u32> {
    let mut num_events = 0;
    raw_event_domain_num_events(event_domain.as_raw(), &mut num_events)?;
    Ok(num_events)
}

pub fn event_domain_enum_events(event_domain: EventDomainId) -> Result<Vec<EventId>> {
    let raw = query_array(|array_size_bytes, event_array| unsafe {
        sys::cuptiEventDomainEnumEvents(event_domain.as_raw(), array_size_bytes, event_array)
    })?;
    Ok(raw.into_iter().map(EventId::from).collect())
}

fn raw_event_get_attribute(
    event: EventId,
    attrib: EventAttribute,
    value_size: *mut u64,
    value: *mut (),
) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiEventGetAttribute(
            event.as_raw(),
            attrib.into(),
            value_size,
            value as _,
        ))?;
    }
    Ok(())
}

pub fn event_attribute_value(
    event: EventId,
    attrib: EventAttribute,
) -> Result<EventAttributeValue> {
    let mut read = |value_size: *mut u64, value: *mut ()| {
        raw_event_get_attribute(event, attrib, value_size, value)
    };
    match attrib {
        EventAttribute::Name => Ok(EventAttributeValue::Name(read_attribute_string(
            &mut read,
            "event_attribute_size",
        )?)),
        EventAttribute::ShortDescription => Ok(EventAttributeValue::ShortDescription(
            read_attribute_string(&mut read, "event_attribute_size")?,
        )),
        EventAttribute::LongDescription => Ok(EventAttributeValue::LongDescription(
            read_attribute_string(&mut read, "event_attribute_size")?,
        )),
        EventAttribute::Category => Ok(EventAttributeValue::Category(EventCategory::from(
            read_attribute_scalar::<sys::CUpti_EventCategory>(&mut read, "event_attribute_size")?,
        ))),
        EventAttribute::ProfilingScope => Ok(EventAttributeValue::ProfilingScope(
            EventProfilingScope::from(read_attribute_scalar::<sys::CUpti_EventProfilingScope>(
                &mut read,
                "event_attribute_size",
            )?),
        )),
        EventAttribute::ForceInt => Err(Error::InvalidAttribute {
            name: format!("{attrib:?}"),
        }),
    }
}

pub fn event_id_from_name(device: Device, event_name: &CStr) -> Result<EventId> {
    let mut event = 0u32;
    unsafe {
        try_ffi!(sys::cuptiEventGetIdFromName(
            device.id() as sys::CUdevice,
            event_name.as_ptr(),
            &mut event,
        ))?;
    }
    Ok(EventId::from(event))
}

pub fn event_id_from_str(device: Device, event_name: &str) -> Result<EventId> {
    let event_name = CString::new(event_name)?;
    event_id_from_name(device, &event_name)
}

pub fn event_group_create(context: &Context) -> Result<EventGroup> {
    let mut event_group = ptr::null_mut();
    unsafe {
        try_ffi!(sys::cuptiEventGroupCreate(
            context.as_raw(),
            &mut event_group,
            0
        ))?;
    }
    EventGroup::from_raw(event_group)
}

unsafe fn raw_event_group_get_attribute(
    event_group: sys::CUpti_EventGroup,
    attrib: EventGroupAttribute,
    value_size: *mut u64,
    value: *mut (),
) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiEventGroupGetAttribute(
            event_group,
            attrib.into(),
            value_size,
            value as _,
        ))?;
    }
    Ok(())
}

fn raw_event_group_set_attribute(
    event_group: sys::CUpti_EventGroup,
    attrib: EventGroupAttribute,
    value_size: u64,
    value: *mut (),
) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiEventGroupSetAttribute(
            event_group,
            attrib.into(),
            value_size,
            value as _,
        ))?;
    }
    Ok(())
}

pub fn event_group_attribute_value(
    event_group: &EventGroup,
    attrib: EventGroupAttribute,
) -> Result<EventGroupAttributeValue> {
    let mut read = |value_size: *mut u64, value: *mut ()| unsafe {
        raw_event_group_get_attribute(event_group.as_raw(), attrib, value_size, value)
    };
    match attrib {
        EventGroupAttribute::EventDomainId => Ok(EventGroupAttributeValue::EventDomainId(
            EventDomainId::from(read_attribute_scalar::<sys::CUpti_EventDomainID>(
                &mut read,
                "event_group_attribute_size",
            )?),
        )),
        EventGroupAttribute::ProfileAllDomainInstances => {
            Ok(EventGroupAttributeValue::ProfileAllDomainInstances(
                read_attribute_scalar::<u32>(&mut read, "event_group_attribute_size")? != 0,
            ))
        }
        EventGroupAttribute::NumEvents => Ok(EventGroupAttributeValue::NumEvents(
            read_attribute_scalar::<u32>(&mut read, "event_group_attribute_size")?,
        )),
        EventGroupAttribute::Events => {
            let raw = query_array::<sys::CUpti_EventID>(|array_size_bytes, event_array| unsafe {
                sys::cuptiEventGroupGetAttribute(
                    event_group.as_raw(),
                    attrib.into(),
                    array_size_bytes,
                    event_array.cast(),
                )
            })?;
            Ok(EventGroupAttributeValue::Events(
                raw.into_iter().map(EventId::from).collect(),
            ))
        }
        EventGroupAttribute::InstanceCount => Ok(EventGroupAttributeValue::InstanceCount(
            read_attribute_scalar::<u32>(&mut read, "event_group_attribute_size")?,
        )),
        EventGroupAttribute::ProfilingScope => Ok(EventGroupAttributeValue::ProfilingScope(
            EventProfilingScope::from(read_attribute_scalar::<sys::CUpti_EventProfilingScope>(
                &mut read,
                "event_group_attribute_size",
            )?),
        )),
        EventGroupAttribute::UserData | EventGroupAttribute::ForceInt => {
            Err(Error::InvalidAttribute {
                name: format!("{attrib:?}"),
            })
        }
    }
}

pub fn event_group_set_attribute_value(
    event_group: &EventGroup,
    attrib: EventGroupAttribute,
    value: EventGroupAttributeSetting,
) -> Result<()> {
    match (attrib, value) {
        (
            EventGroupAttribute::ProfileAllDomainInstances,
            EventGroupAttributeSetting::ProfileAllDomainInstances(value),
        ) => {
            let mut value = value as u32;
            let value_size = to_u64(size_of::<u32>(), "event_group_attribute_size")?;
            raw_event_group_set_attribute(
                event_group.as_raw(),
                attrib,
                value_size,
                (&mut value as *mut u32).cast(),
            )
        }
        _ => Err(Error::InvalidAttribute {
            name: format!("{attrib:?}"),
        }),
    }
}

pub fn event_group_add_event(event_group: &EventGroup, event: EventId) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiEventGroupAddEvent(
            event_group.as_raw(),
            event.as_raw(),
        ))?;
    }
    Ok(())
}

pub fn event_group_remove_event(event_group: &EventGroup, event: EventId) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiEventGroupRemoveEvent(
            event_group.as_raw(),
            event.as_raw(),
        ))?;
    }
    Ok(())
}

pub fn event_group_remove_all_events(event_group: &EventGroup) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiEventGroupRemoveAllEvents(event_group.as_raw()))?;
    }
    Ok(())
}

pub fn event_group_reset_all_events(event_group: &EventGroup) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiEventGroupResetAllEvents(event_group.as_raw()))?;
    }
    Ok(())
}

pub fn event_group_enable(event_group: &EventGroup) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiEventGroupEnable(event_group.as_raw()))?;
    }
    Ok(())
}

pub fn event_group_disable(event_group: &EventGroup) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiEventGroupDisable(event_group.as_raw()))?;
    }
    Ok(())
}

fn raw_event_group_read_event(
    event_group: sys::CUpti_EventGroup,
    flags: ReadEventFlags,
    event: EventId,
    event_value_buffer_size_bytes: *mut u64,
    event_value_buffer: *mut u64,
) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiEventGroupReadEvent(
            event_group,
            flags.into(),
            event.as_raw(),
            event_value_buffer_size_bytes,
            event_value_buffer,
        ))?;
    }
    Ok(())
}

pub fn event_group_read_event(
    event_group: &EventGroup,
    flags: ReadEventFlags,
    event: EventId,
) -> Result<Vec<u64>> {
    let mut value_size_bytes = 0;
    raw_event_group_read_event(
        event_group.as_raw(),
        flags,
        event,
        &mut value_size_bytes,
        ptr::null_mut(),
    )?;
    let len = to_usize(value_size_bytes, "event_value_buffer_size_bytes")? / size_of::<u64>();
    let mut values = vec![0; len];
    raw_event_group_read_event(
        event_group.as_raw(),
        flags,
        event,
        &mut value_size_bytes,
        values.as_mut_ptr(),
    )?;
    let len = to_usize(value_size_bytes, "event_value_buffer_size_bytes")? / size_of::<u64>();
    values.truncate(len);
    Ok(values)
}

fn raw_event_group_read_all_events(
    event_group: sys::CUpti_EventGroup,
    flags: ReadEventFlags,
    event_value_buffer_size_bytes: *mut u64,
    event_value_buffer: *mut u64,
    event_id_array_size_bytes: *mut u64,
    event_id_array: *mut sys::CUpti_EventID,
    num_event_ids_read: *mut u64,
) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiEventGroupReadAllEvents(
            event_group,
            flags.into(),
            event_value_buffer_size_bytes,
            event_value_buffer,
            event_id_array_size_bytes,
            event_id_array,
            num_event_ids_read,
        ))?;
    }
    Ok(())
}

pub fn event_group_read_all_events(
    event_group: &EventGroup,
    flags: ReadEventFlags,
) -> Result<EventGroupReadAll> {
    let mut value_size_bytes = 0;
    let mut event_id_size_bytes = 0;
    let mut num_event_ids_read = 0;
    raw_event_group_read_all_events(
        event_group.as_raw(),
        flags,
        &mut value_size_bytes,
        ptr::null_mut(),
        &mut event_id_size_bytes,
        ptr::null_mut(),
        &mut num_event_ids_read,
    )?;
    let value_len = to_usize(value_size_bytes, "event_value_buffer_size_bytes")? / size_of::<u64>();
    let event_id_len = to_usize(event_id_size_bytes, "event_id_array_size_bytes")?
        / size_of::<sys::CUpti_EventID>();
    let mut values = vec![0; value_len];
    let mut event_ids = vec![0; event_id_len];
    raw_event_group_read_all_events(
        event_group.as_raw(),
        flags,
        &mut value_size_bytes,
        values.as_mut_ptr(),
        &mut event_id_size_bytes,
        event_ids.as_mut_ptr(),
        &mut num_event_ids_read,
    )?;
    let value_len = to_usize(value_size_bytes, "event_value_buffer_size_bytes")? / size_of::<u64>();
    let event_id_len = to_usize(num_event_ids_read, "num_event_ids_read")?;
    values.truncate(value_len);
    event_ids.truncate(event_id_len);
    Ok(EventGroupReadAll {
        event_ids: event_ids.into_iter().map(EventId::from).collect(),
        values,
    })
}

pub fn event_group_sets_create(context: &Context, event_ids: &[EventId]) -> Result<EventGroupSets> {
    let mut event_group_sets = ptr::null_mut();
    let event_ids_size = checked_byte_len::<EventId>(event_ids.len(), "event_ids_size_bytes")?;
    let mut event_ids = event_ids
        .iter()
        .copied()
        .map(EventId::as_raw)
        .collect::<Vec<_>>();
    unsafe {
        try_ffi!(sys::cuptiEventGroupSetsCreate(
            context.as_raw(),
            to_u64(event_ids_size, "event_ids_size_bytes")?,
            event_ids.as_mut_ptr(),
            &mut event_group_sets,
        ))?;
    }
    EventGroupSets::from_raw(event_group_sets)
}

pub fn enable_kernel_replay_mode(context: &Context) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiEnableKernelReplayMode(context.as_raw()))?;
    }
    Ok(())
}

pub fn disable_kernel_replay_mode(context: &Context) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiDisableKernelReplayMode(context.as_raw()))?;
    }
    Ok(())
}

pub fn kernel_replay_subscribe_update<F>(callback: F) -> Result<()>
where
    F: FnMut(&CStr, i32) + Send + 'static,
{
    *KERNEL_REPLAY_UPDATE_CALLBACK
        .lock()
        .map_err(|_| Error::LockPoisoned {
            name: "kernel replay update callback".into(),
        })? = Some(Box::new(callback));
    unsafe {
        try_ffi!(sys::cuptiKernelReplaySubscribeUpdate(
            Some(kernel_replay_update_trampoline),
            ptr::null_mut(),
        ))?;
    }
    Ok(())
}

unsafe extern "C" fn kernel_replay_update_trampoline(
    kernel_name: *const c_char,
    num_replays_done: i32,
    _custom_data: *mut core::ffi::c_void,
) {
    if kernel_name.is_null() {
        return;
    }

    let _ = catch_unwind(AssertUnwindSafe(|| {
        if let Ok(mut callback) = KERNEL_REPLAY_UPDATE_CALLBACK.lock()
            && let Some(callback) = callback.as_mut()
        {
            callback(unsafe { CStr::from_ptr(kernel_name) }, num_replays_done);
        }
    }));
}

fn raw_num_metrics(num_metrics: *mut u32) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiGetNumMetrics(num_metrics))?;
    }
    Ok(())
}

pub fn num_metrics() -> Result<u32> {
    let mut num_metrics = 0;
    raw_num_metrics(&mut num_metrics)?;
    Ok(num_metrics)
}

pub fn enum_metrics() -> Result<Vec<MetricId>> {
    let raw = query_array(|array_size_bytes, metric_array| unsafe {
        sys::cuptiEnumMetrics(array_size_bytes, metric_array)
    })?;
    Ok(raw.into_iter().map(MetricId::from).collect())
}

fn raw_device_num_metrics(device: sys::CUdevice, num_metrics: *mut u32) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiDeviceGetNumMetrics(device, num_metrics))?;
    }
    Ok(())
}

pub fn device_num_metrics(device: Device) -> Result<u32> {
    let mut num_metrics = 0;
    raw_device_num_metrics(device.id() as sys::CUdevice, &mut num_metrics)?;
    Ok(num_metrics)
}

pub fn device_enum_metrics(device: Device) -> Result<Vec<MetricId>> {
    let raw = query_array(|array_size_bytes, metric_array| unsafe {
        sys::cuptiDeviceEnumMetrics(device.id() as sys::CUdevice, array_size_bytes, metric_array)
    })?;
    Ok(raw.into_iter().map(MetricId::from).collect())
}

fn raw_metric_get_attribute(
    metric: MetricId,
    attrib: MetricAttribute,
    value_size: *mut u64,
    value: *mut (),
) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiMetricGetAttribute(
            metric.as_raw(),
            attrib.into(),
            value_size,
            value as _,
        ))?;
    }
    Ok(())
}

pub fn metric_attribute_value(
    metric: MetricId,
    attrib: MetricAttribute,
) -> Result<MetricAttributeValue> {
    let mut read = |value_size: *mut u64, value: *mut ()| {
        raw_metric_get_attribute(metric, attrib, value_size, value)
    };
    match attrib {
        MetricAttribute::Name => Ok(MetricAttributeValue::Name(read_attribute_string(
            &mut read,
            "metric_attribute_size",
        )?)),
        MetricAttribute::ShortDescription => Ok(MetricAttributeValue::ShortDescription(
            read_attribute_string(&mut read, "metric_attribute_size")?,
        )),
        MetricAttribute::LongDescription => Ok(MetricAttributeValue::LongDescription(
            read_attribute_string(&mut read, "metric_attribute_size")?,
        )),
        MetricAttribute::Category => Ok(MetricAttributeValue::Category(MetricCategory::from(
            read_attribute_scalar::<sys::CUpti_MetricCategory>(&mut read, "metric_attribute_size")?,
        ))),
        MetricAttribute::ValueKind => Ok(MetricAttributeValue::ValueKind(MetricValueKind::from(
            read_attribute_scalar::<sys::CUpti_MetricValueKind>(
                &mut read,
                "metric_attribute_size",
            )?,
        ))),
        MetricAttribute::EvaluationMode => Ok(MetricAttributeValue::EvaluationMode(
            MetricEvaluationMode::from(read_attribute_scalar::<sys::CUpti_MetricEvaluationMode>(
                &mut read,
                "metric_attribute_size",
            )?),
        )),
        MetricAttribute::ForceInt => Err(Error::InvalidAttribute {
            name: format!("{attrib:?}"),
        }),
    }
}

pub fn metric_id_from_name(device: Device, metric_name: &CStr) -> Result<MetricId> {
    let mut metric = 0u32;
    unsafe {
        try_ffi!(sys::cuptiMetricGetIdFromName(
            device.id() as sys::CUdevice,
            metric_name.as_ptr(),
            &mut metric,
        ))?;
    }
    Ok(MetricId::from(metric))
}

pub fn metric_id_from_str(device: Device, metric_name: &str) -> Result<MetricId> {
    let metric_name = CString::new(metric_name)?;
    metric_id_from_name(device, &metric_name)
}

fn raw_metric_num_events(metric: MetricId, num_events: *mut u32) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiMetricGetNumEvents(metric.as_raw(), num_events))?;
    }
    Ok(())
}

pub fn metric_num_events(metric: MetricId) -> Result<u32> {
    let mut num_events = 0;
    raw_metric_num_events(metric, &mut num_events)?;
    Ok(num_events)
}

pub fn metric_enum_events(metric: MetricId) -> Result<Vec<EventId>> {
    let raw = query_array(|array_size_bytes, event_array| unsafe {
        sys::cuptiMetricEnumEvents(metric.as_raw(), array_size_bytes, event_array)
    })?;
    Ok(raw.into_iter().map(EventId::from).collect())
}

fn raw_metric_num_properties(metric: MetricId, num_properties: *mut u32) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiMetricGetNumProperties(
            metric.as_raw(),
            num_properties
        ))?;
    }
    Ok(())
}

pub fn metric_num_properties(metric: MetricId) -> Result<u32> {
    let mut num_properties = 0;
    raw_metric_num_properties(metric, &mut num_properties)?;
    Ok(num_properties)
}

pub fn metric_enum_properties(metric: MetricId) -> Result<Vec<MetricPropertyId>> {
    let raw = query_array(|array_size_bytes, property_array| unsafe {
        sys::cuptiMetricEnumProperties(metric.as_raw(), array_size_bytes, property_array)
    })?;
    Ok(raw.into_iter().map(MetricPropertyId::from).collect())
}

fn raw_metric_required_event_group_sets(
    context: sys::CUcontext,
    metric: MetricId,
    event_group_sets: *mut *mut sys::CUpti_EventGroupSets,
) -> Result<()> {
    unsafe {
        try_ffi!(sys::cuptiMetricGetRequiredEventGroupSets(
            context,
            metric.as_raw(),
            event_group_sets,
        ))?;
    }
    Ok(())
}

pub fn metric_required_event_group_sets(
    context: &Context,
    metric: MetricId,
) -> Result<EventGroupSets> {
    let mut event_group_sets = ptr::null_mut();
    raw_metric_required_event_group_sets(context.as_raw(), metric, &mut event_group_sets)?;
    EventGroupSets::from_raw(event_group_sets)
}

pub fn metric_create_event_group_sets(
    context: &Context,
    metric_ids: &[MetricId],
) -> Result<EventGroupSets> {
    let mut event_group_sets = ptr::null_mut();
    let metric_ids_size = checked_byte_len::<MetricId>(metric_ids.len(), "metric_ids_size_bytes")?;
    let mut metric_ids = metric_ids
        .iter()
        .copied()
        .map(MetricId::as_raw)
        .collect::<Vec<_>>();
    unsafe {
        try_ffi!(sys::cuptiMetricCreateEventGroupSets(
            context.as_raw(),
            to_u64(metric_ids_size, "metric_ids_size_bytes")?,
            metric_ids.as_mut_ptr(),
            &mut event_group_sets,
        ))?;
    }
    EventGroupSets::from_raw(event_group_sets)
}

pub fn metric_value(
    device: Device,
    metric: MetricId,
    event_ids: &[EventId],
    event_values: &[u64],
    duration: u64,
) -> Result<MetricValue> {
    let event_ids_size = checked_byte_len::<EventId>(event_ids.len(), "event_ids_size_bytes")?;
    let event_values_size = checked_byte_len::<u64>(event_values.len(), "event_values_size_bytes")?;
    let mut metric_value = sys::CUpti_MetricValue::default();
    let mut event_ids = event_ids
        .iter()
        .copied()
        .map(EventId::as_raw)
        .collect::<Vec<_>>();
    let mut event_values = event_values.to_vec();
    unsafe {
        try_ffi!(sys::cuptiMetricGetValue(
            device.id() as sys::CUdevice,
            metric.as_raw(),
            to_u64(event_ids_size, "event_ids_size_bytes")?,
            event_ids.as_mut_ptr(),
            to_u64(event_values_size, "event_values_size_bytes")?,
            event_values.as_mut_ptr(),
            duration,
            &mut metric_value,
        ))?;
    }
    metric_value_from_raw(metric, metric_value)
}

pub fn metric_value2(
    metric: MetricId,
    event_ids: &[EventId],
    event_values: &[u64],
    property_ids: &[MetricPropertyId],
    property_values: &[u64],
) -> Result<MetricValue> {
    let event_ids_size = checked_byte_len::<EventId>(event_ids.len(), "event_ids_size_bytes")?;
    let event_values_size = checked_byte_len::<u64>(event_values.len(), "event_values_size_bytes")?;
    let property_ids_size =
        checked_byte_len::<MetricPropertyId>(property_ids.len(), "property_ids_size_bytes")?;
    let property_values_size =
        checked_byte_len::<u64>(property_values.len(), "property_values_size_bytes")?;
    let mut property_ids = property_ids
        .iter()
        .copied()
        .map(sys::CUpti_MetricPropertyID::from)
        .collect::<Vec<_>>();
    let mut event_ids = event_ids
        .iter()
        .copied()
        .map(EventId::as_raw)
        .collect::<Vec<_>>();
    let mut event_values = event_values.to_vec();
    let mut property_values = property_values.to_vec();
    let mut metric_value = sys::CUpti_MetricValue::default();
    unsafe {
        try_ffi!(sys::cuptiMetricGetValue2(
            metric.as_raw(),
            to_u64(event_ids_size, "event_ids_size_bytes")?,
            event_ids.as_mut_ptr(),
            to_u64(event_values_size, "event_values_size_bytes")?,
            event_values.as_mut_ptr(),
            to_u64(property_ids_size, "property_ids_size_bytes")?,
            property_ids.as_mut_ptr(),
            to_u64(property_values_size, "property_values_size_bytes")?,
            property_values.as_mut_ptr(),
            &mut metric_value,
        ))?;
    }
    metric_value_from_raw(metric, metric_value)
}

fn metric_value_from_raw(metric: MetricId, value: sys::CUpti_MetricValue) -> Result<MetricValue> {
    let MetricAttributeValue::ValueKind(kind) =
        metric_attribute_value(metric, MetricAttribute::ValueKind)?
    else {
        return Err(Error::InvalidAttribute {
            name: format!("{:?}", MetricAttribute::ValueKind),
        });
    };
    let value = unsafe {
        match kind {
            MetricValueKind::Double => MetricValue::Double(value.metricValueDouble),
            MetricValueKind::Uint64 => MetricValue::Uint64(value.metricValueUint64),
            MetricValueKind::Percent => MetricValue::Percent(value.metricValuePercent),
            MetricValueKind::Throughput => MetricValue::Throughput(value.metricValueThroughput),
            MetricValueKind::Int64 => MetricValue::Int64(value.metricValueInt64),
            MetricValueKind::UtilizationLevel => {
                MetricValue::UtilizationLevel(value.metricValueUtilizationLevel.into())
            }
            MetricValueKind::NvtxExtendedPayload => {
                MetricValue::NvtxExtendedPayload(value.metricValueNvtxExtendedPayload)
            }
            MetricValueKind::ForceInt => MetricValue::Uint64(value.metricValueUint64),
        }
    };
    Ok(value)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn string_attribute_reader_trims_at_nul() -> Result<()> {
        let value = read_attribute_string(
            |value_size, value| {
                let bytes = b"metric-name\0ignored";
                unsafe {
                    bytes
                        .as_ptr()
                        .copy_to_nonoverlapping(value.cast::<u8>(), bytes.len());
                    *value_size = bytes.len() as u64;
                }
                Ok(())
            },
            "test_attribute_size",
        )?;

        assert_eq!(value, "metric-name");
        Ok(())
    }

    #[test]
    fn event_domain_attribute_reader_maps_collection_method() -> Result<()> {
        let value = event_domain_attribute_value_with(
            &mut |value_size, value| {
                let mut raw = sys::CUpti_EventCollectionMethod::CUPTI_EVENT_COLLECTION_METHOD_SM;
                unsafe {
                    *value_size = std::mem::size_of_val(&raw) as u64;
                    std::ptr::copy_nonoverlapping(&mut raw as *mut _, value.cast(), 1);
                }
                Ok(())
            },
            EventDomainAttribute::CollectionMethod,
        )?;

        assert_eq!(
            value,
            EventDomainAttributeValue::CollectionMethod(EventCollectionMethod::Sm)
        );
        Ok(())
    }

    #[test]
    fn event_group_setter_rejects_unsupported_attributes_before_ffi() {
        let event_group = EventGroup {
            handle: std::ptr::NonNull::<std::ffi::c_void>::dangling().as_ptr(),
        };
        let error = event_group_set_attribute_value(
            &event_group,
            EventGroupAttribute::UserData,
            EventGroupAttributeSetting::ProfileAllDomainInstances(true),
        )
        .unwrap_err();

        std::mem::forget(event_group);
        assert!(matches!(error, Error::InvalidAttribute { .. }));
    }
}