use std::sync::Arc;

use cudarc::driver::{CudaContext, CudaSlice, CudaStream, DevicePtrMut, DriverError};
use nvtiff_sys::{
    nvtiffDecodeCheckSupported, nvtiffDecodeImage, nvtiffDecodeParams, nvtiffDecoder,
    nvtiffDecoderCreateSimple, nvtiffStatus_t, nvtiffStream, nvtiffStreamCreate,
    nvtiffStreamParseFromFile,
};

#[test]
fn test_decode_image() -> Result<(), DriverError> {
    // Step 1: Create stream pointers/handles on host (CPU) and device (GPU)

    // Step 1a: Init TIFF stream on host (CPU)
    let mut host_stream = std::mem::MaybeUninit::uninit();
    let tiff_stream: *mut *mut nvtiffStream = host_stream.as_mut_ptr();

    let status_cpustream: nvtiffStatus_t::Type = unsafe { nvtiffStreamCreate(tiff_stream) };
    assert_eq!(status_cpustream, nvtiffStatus_t::NVTIFF_STATUS_SUCCESS);

    // Step 1b: Init CUDA stream on device (GPU)
    let ctx: Arc<CudaContext> = cudarc::driver::CudaContext::new(0)?; // Set on GPU:0
    let stream: Arc<CudaStream> = ctx.default_stream();
    let cuda_stream: *mut nvtiff_sys::CUstream_st = stream.cu_stream().cast::<_>();

    // Step 1c: Init decoder handle
    let mut decoder_handle = std::mem::MaybeUninit::uninit();
    let nvtiff_decoder: *mut *mut nvtiffDecoder = decoder_handle.as_mut_ptr();

    let status_decoder: u32 = unsafe { nvtiffDecoderCreateSimple(nvtiff_decoder, cuda_stream) };
    dbg!(status_decoder);
    assert_eq!(status_decoder, nvtiffStatus_t::NVTIFF_STATUS_SUCCESS);

    // Step 2: Parse the TIFF file from disk
    let tiff_cstr = std::ffi::CString::new("images/float32.tif").expect("CString::new failed");
    let tiff_path: *const std::os::raw::c_char = tiff_cstr.as_ptr();

    let status_parse: u32 = unsafe { nvtiffStreamParseFromFile(tiff_path, *tiff_stream) };
    dbg!(status_parse);
    assert_eq!(status_parse, nvtiffStatus_t::NVTIFF_STATUS_SUCCESS);

    // Step 3a: Check if image is supported first
    let mut params = std::mem::MaybeUninit::uninit();
    let decode_params: *mut nvtiffDecodeParams = params.as_mut_ptr();
    let status_check: u32 = unsafe {
        nvtiffDecodeCheckSupported(
            *tiff_stream,
            *nvtiff_decoder,
            decode_params,
            0, //image_id
        )
    };
    dbg!(status_check);
    assert_eq!(status_check, nvtiffStatus_t::NVTIFF_STATUS_SUCCESS);

    // Step 3b: Allocate memory on device, get pointer, do the TIFF decoding
    let image_size: usize = 3 * 2 * 4; // Width:3, Height:2, 4 bytes per f32 num
    let mut image_stream: CudaSlice<u8> = stream.alloc_zeros::<u8>(image_size)?;
    let image_ptr: u64 = image_stream.device_ptr_mut(&stream).0;
    let image_out_d = image_ptr as *mut std::ffi::c_void;
    let status_decode: u32 = unsafe {
        nvtiffDecodeImage(
            *tiff_stream,
            *nvtiff_decoder,
            decode_params,
            0, // image_id
            image_out_d,
            cuda_stream,
        )
    };
    dbg!(status_decode); // 2: NVTIFF_STATUS_INVALID_PARAMETER; 8: NVTIFF_STATUS_INTERNAL_ERROR
    assert_eq!(status_decode, nvtiffStatus_t::NVTIFF_STATUS_SUCCESS);

    // Step 2c: Transfer decoded bytes from device to host, and check results
    let mut image_out_h: Vec<u8> = vec![0; image_size];
    stream.memcpy_dtoh(&image_stream, &mut image_out_h)?;
    dbg!(image_out_h.clone());
    let float_array: Vec<f32> = image_out_h
        // https://stackoverflow.com/questions/77388769/convert-vecu8-to-vecfloat-in-rust
        // .array_chunks::<4>()
        // .copied()
        .chunks_exact(4)
        .map(TryInto::try_into)
        .map(Result::unwrap)
        .map(f32::from_le_bytes)
        .collect();
    assert_eq!(float_array, vec![1.41, 1.23, 0.78, 0.32, -0.23, -1.88]);

    Ok(())
}
