use std::ffi::c_void;
use xserv_tensor::{DType, Device, Tensor};

unsafe extern "C" {
    fn launch_argmax_bf16(
        logits: *const c_void,
        out_idx: *mut c_void,
        rows: i32,
        cols: i32,
        stream: *mut c_void,
    );
}

/// GPU argmax over the last dim of a [rows, cols] BF16 tensor.
///
/// Returns a host `Vec<u32>` of length `rows`. Internally:
///   - launches one kernel that writes [rows] i32 indices on device
///   - D2H copies just `rows * 4` bytes (vs `rows * cols * 2` for the
///     "copy logits to CPU then argmax" path it replaces)
///
/// This is the greedy-decode hot path: avoids touching the full
/// [B, vocab] logits buffer on the host every step.
pub fn argmax_bf16_to_host(logits: &Tensor) -> Vec<u32> {
    assert_eq!(logits.ndim(), 2, "argmax expects a 2D [rows, cols] tensor");
    assert_eq!(logits.dtype(), DType::BF16, "argmax kernel is BF16-only");
    assert!(logits.is_contiguous(), "argmax requires contiguous input");
    assert!(
        matches!(logits.device(), Device::Cuda(_)),
        "argmax requires GPU input"
    );

    let rows = logits.shape()[0];
    let cols = logits.shape()[1];
    assert!(rows <= i32::MAX as usize);
    assert!(cols <= i32::MAX as usize);

    // Output buffer: rows * i32. Pooled allocator so this is essentially free
    // after the first call.
    let bytes = rows * std::mem::size_of::<i32>();
    let mut out = xserv_cuda::allocator::cached_alloc(bytes).expect("argmax out alloc");

    unsafe {
        launch_argmax_bf16(
            logits.data_ptr() as *const c_void,
            out.as_mut_ptr() as *mut c_void,
            rows as i32,
            cols as i32,
            xserv_cuda::current_stream_raw(),
        );
    }

    let mut host_bytes = vec![0u8; bytes];
    out.copy_to_host(&mut host_bytes).expect("argmax D2H");
    drop(out); // returned to pool

    let host_i32: &[i32] =
        unsafe { std::slice::from_raw_parts(host_bytes.as_ptr() as *const i32, rows) };
    host_i32.iter().map(|&v| v as u32).collect()
}

/// Convenience: argmax of a single row [1, cols] (or [cols] reshaped to [1, cols]).
pub fn argmax_bf16_single(logits: &Tensor) -> u32 {
    let cols = *logits.shape().last().unwrap();
    let rows = logits.numel() / cols;
    assert_eq!(rows, 1, "argmax_bf16_single requires a single row");
    let view = if logits.ndim() == 2 {
        logits.clone()
    } else {
        logits.reshape(&[1, cols])
    };
    argmax_bf16_to_host(&view)[0]
}