use pyo3::prelude::*;
use std::ffi::c_int;

// Link C++ functions
#[link(name = "trme_core", kind = "static")]
extern "C" {
    fn rsr_gemm_dispatch(
        A: *const i8,
        B: *const i8,
        C: *mut i32,
        bias: *const i32,
        N: c_int,
        K: c_int,
        M: c_int,
        block_size: c_int,
    );
}

// Python Wrapper
#[pyfunction]
fn rsr_gemm(
    py: Python,
    a_bytes: &[u8],
    b_bytes: &[u8],
    n: i32,
    k: i32,
    m: i32,
    block_size: i32,
) -> PyResult<Vec<i32>> {
    // In a real optimized binding, we would use numpy arrays directly via PyArray.
    // For this prototype, we accept bytes/lists and return a vec.

    let mut c_out = vec![0i32; (n * m) as usize];

    // Unsafe call to C++ kernel
    unsafe {
        rsr_gemm_dispatch(
            a_bytes.as_ptr() as *const i8,
            b_bytes.as_ptr() as *const i8,
            c_out.as_mut_ptr(),
            std::ptr::null(), // No bias for basic binding
            n,
            k,
            m,
            block_size,
        );
    }

    Ok(c_out)
}

#[pymodule]
fn _trme_backend(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(rsr_gemm, m)?)?;
    Ok(())
}
