use crate::error::{self, Result};
use crate::ffi;
use std::ffi::CStr;
use std::os::raw::c_char;

#[derive(Debug, Clone)]
pub struct DeviceInfo {
    pub index: u32,
    pub name: String,
    pub total_memory: usize,
    pub free_memory: usize,
    pub compute_major: i32,
    pub compute_minor: i32,
    pub sm_count: i32,
    pub shared_mem_per_block: usize,
    pub warp_size: i32,
    pub max_threads_per_block: i32,
}

unsafe extern "C" {
    fn cudaDeviceGetAttribute(value: *mut i32, attr: i32, device: i32) -> i32;
    fn cudaMemGetInfo(free: *mut usize, total: *mut usize) -> i32;
}

fn get_attr(attr: i32, device: u32) -> Result<i32> {
    let mut value = 0;
    error::check(unsafe { cudaDeviceGetAttribute(&mut value, attr, device as i32) })?;
    Ok(value)
}

pub fn device_count() -> Result<i32> {
    let mut count = 0;
    error::check(unsafe { ffi::cudaGetDeviceCount(&mut count) })?;
    Ok(count)
}

pub fn set_device(device: u32) -> Result<()> {
    error::check(unsafe { ffi::cudaSetDevice(device as i32) })
}

pub fn current_device() -> Result<u32> {
    let mut dev = 0;
    error::check(unsafe { ffi::cudaGetDevice(&mut dev) })?;
    Ok(dev as u32)
}

pub fn device_info(device: u32) -> Result<DeviceInfo> {
    // Heap-allocate oversized buffer for cudaDeviceProp (layout varies by CUDA version).
    // CUDA 12.x struct is ~5-6 KB; use 32 KB to guard against future growth.
    let mut prop_buf = vec![0u8; 32768];
    error::check(unsafe { ffi::cudaGetDeviceProperties(prop_buf.as_mut_ptr(), device as i32) })?;
    // Name is always the first field: char[256].
    let name = unsafe { CStr::from_ptr(prop_buf.as_ptr() as *const c_char) }
        .to_string_lossy()
        .into_owned();

    // Get memory info via cudaMemGetInfo (layout-independent).
    let prev = current_device()?;
    set_device(device)?;
    let mut free = 0usize;
    let mut total = 0usize;
    error::check(unsafe { cudaMemGetInfo(&mut free, &mut total) })?;
    if prev != device {
        set_device(prev)?;
    }

    // Attribute IDs from cuda_runtime_api.h
    const SHARED_MEM_PER_BLOCK: i32 = 8;
    const WARP_SIZE: i32 = 10;
    const MAX_THREADS_PER_BLOCK: i32 = 1;
    const MULTI_PROCESSOR_COUNT: i32 = 16;
    const COMPUTE_MAJOR: i32 = 75;
    const COMPUTE_MINOR: i32 = 76;

    Ok(DeviceInfo {
        index: device,
        name,
        total_memory: total,
        free_memory: free,
        compute_major: get_attr(COMPUTE_MAJOR, device)?,
        compute_minor: get_attr(COMPUTE_MINOR, device)?,
        sm_count: get_attr(MULTI_PROCESSOR_COUNT, device)?,
        shared_mem_per_block: get_attr(SHARED_MEM_PER_BLOCK, device)? as usize,
        warp_size: get_attr(WARP_SIZE, device)?,
        max_threads_per_block: get_attr(MAX_THREADS_PER_BLOCK, device)?,
    })
}

pub fn synchronize() -> Result<()> {
    error::check(unsafe { ffi::cudaDeviceSynchronize() })
}