use xrt::ffi::XCL_BO_FLAGS_NONE;
use xrt::native::buffer::SyncDirection;
use xrt::native::buffer::XRTBuffer;
use xrt::native::device::XRTDevice;
use xrt::native::kernel::XRTKernel;
use xrt::native::run::ERTCommandState;
use xrt::native::run::XRTRun;
use xrt::utils::get_xclbin_path;
use xrt::Result;

mod data;

use data::{VScaleTestData, SIZE};

fn run_vscale_native<T: VScaleTestData + std::fmt::Debug + Copy + std::cmp::PartialEq<T>>(
) -> Result<()> {
    std::env::set_var("XCL_EMULATION_MODE", "sw_emu");

    let mut device = XRTDevice::try_from(0)?;
    device.load_xclbin(get_xclbin_path(&format!("./hls/vscale_{}", T::name())).as_str())?;
    let add_kernel = XRTKernel::new(&format!("vscale_{}", T::name()), &device)?;
    let add_run = XRTRun::try_from(&add_kernel)?;
    let in_buffer = XRTBuffer::new(
        &device,
        SIZE * std::mem::size_of::<T>(),
        XCL_BO_FLAGS_NONE,
        add_kernel.get_memory_group_for_argument(2)?,
    )?;
    let out_buffer = XRTBuffer::new(
        &device,
        SIZE * std::mem::size_of::<T>(),
        XCL_BO_FLAGS_NONE,
        add_kernel.get_memory_group_for_argument(3)?,
    )?;

    let input: [T; SIZE] = [T::input(); SIZE];
    in_buffer.write(&input, 0)?;
    in_buffer.sync::<T>(SyncDirection::HostToDevice, None, 0)?;

    // Set args
    add_run.set_scalar_argument(0, SIZE)?;
    add_run.set_scalar_argument(1, T::scale())?;
    add_run.set_buffer_argument(2, &in_buffer)?;
    add_run.set_buffer_argument(3, &out_buffer)?;

    // Run
    let _start_state = add_run.start()?;

    let result_state = add_run.wait_for(1000)?;
    assert_eq!(result_state, ERTCommandState::Completed);

    // Get back data
    let mut output: [T; SIZE] = [T::zero(); SIZE];
    out_buffer.sync::<T>(SyncDirection::DeviceToHost, None, 0)?;
    out_buffer.read(&mut output, 0)?;

    // Check result
    for elem in output {
        assert_eq!(elem, T::output());
    }
    Ok(())
}

#[test]
fn run_vscale_native_u32() -> Result<()> {
    run_vscale_native::<u32>()
}

#[test]
fn run_vscale_native_i32() -> Result<()> {
    run_vscale_native::<i32>()
}

#[test]
fn run_vscale_native_u64() -> Result<()> {
    run_vscale_native::<u64>()
}

#[test]
fn run_vscale_native_i64() -> Result<()> {
    run_vscale_native::<i64>()
}

#[test]
fn run_vscale_native_f32() -> Result<()> {
    run_vscale_native::<f32>()
}

#[test]
fn run_vscale_native_f64() -> Result<()> {
    run_vscale_native::<f64>()
}