use lonsdaleite::*;

fn main() {
    let devices: Vec<Device> = Device::enumerate();
    let device: Device = devices.into_iter().find_map(|d| Some(d)).unwrap();

    let command_queue: CommandQueue = device.new_command_queue();

    let data = [0u32, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];

    let input: Buffer = device.new_buffer(
        &BufferDescriptor {
            cache_mode: CPUCacheMode::Default,
            storage_mode: StorageMode::Shared,
        },
        &data,
    );
    let result: Buffer = device.new_buffer_with_size(
        &BufferDescriptor {
            cache_mode: CPUCacheMode::Default,
            storage_mode: StorageMode::Shared,
        },
        1,
    );

    let library: Library =
        device.new_library(&LibraryDescriptor::Source(include_str!("sum.metal")));
    let kernel: Function = library.get_function("sum");

    let compute_pipeline: ComputePipeline = device.new_compute_pipeline(&kernel);

    let command_buffer: CommandBuffer = command_queue.new_command_buffer();
    let compute_encoder: ComputeCommandEncoder =
        command_buffer.new_compute_command_encoder(CommandBufferDispatchType::Concurrent);

    compute_encoder.set_buffer(0, 0, &input);
    compute_encoder.set_buffer(1, 0, &result);
    compute_encoder.set_compute_pipeline(&compute_pipeline);

    compute_encoder.dispatch_thread_groups(
        Size3D {
            width: 1,
            height: 1,
            depth: 1,
        },
        Size3D {
            width: data.len() as u64,
            height: 1,
            depth: 1,
        },
    );
    compute_encoder.end_compute_encoding();
    command_buffer.commit();
    command_buffer.wait_until_completed();

    unsafe {
        println!("{}", *(result.get_raw_contents_as::<u32>()));
    }
}