//! Standalone cublasLtMatmul debug test.
//! Tests the EXACT same API call that fails in the experience collector on H100.
//! Run: SQLX_OFFLINE=true cargo test -p ml --lib -- cuda_pipeline::cublaslt_debug --nocapture

#[cfg(test)]
mod tests {
    use cudarc::cublaslt::sys as cublaslt_sys;
    use cudarc::cublaslt::result as cublaslt_result;
    use cudarc::driver::CudaContext;

    struct TestCase {
        m: i32,
        n: i32,
        k: i32,
        label: &'static str,
    }

    fn test_cublaslt_matmul(
        lt_handle: cublaslt_sys::cublasLtHandle_t,
        stream: cudarc::driver::sys::CUstream,
        ws_ptr: u64,
        ws_size: usize,
        tc: &TestCase,
    ) -> Result<(), String> {
        let m = tc.m as usize;
        let n = tc.n as usize;
        let k = tc.k as usize;

        // Allocate buffers via cuMemAlloc (NOT stream-ordered)
        let mut d_w: cudarc::driver::sys::CUdeviceptr = 0;
        let mut d_a: cudarc::driver::sys::CUdeviceptr = 0;
        let mut d_c: cudarc::driver::sys::CUdeviceptr = 0;
        unsafe {
            cudarc::driver::sys::cuMemAlloc_v2(&mut d_w, k * m * 4);
            cudarc::driver::sys::cuMemAlloc_v2(&mut d_a, k * n * 4);
            cudarc::driver::sys::cuMemAlloc_v2(&mut d_c, m * n * 4);
            cudarc::driver::sys::cuMemsetD8_v2(d_w, 0, k * m * 4);
            cudarc::driver::sys::cuMemsetD8_v2(d_a, 0, k * n * 4);
            cudarc::driver::sys::cuMemsetD8_v2(d_c, 0, m * n * 4);
        }

        let alpha = 1.0_f32;
        let beta = 0.0_f32;
        let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F;
        let compute_type = cublaslt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F;

        let result = unsafe {
            // Create matmul descriptor
            let matmul_desc = cublaslt_result::create_matmul_desc(compute_type, f32_type)
                .map_err(|e| format!("create_matmul_desc: {e:?}"))?;

            // TRANSA=T, TRANSB=N
            let transa: i32 = 1;
            let transb: i32 = 0;
            cublaslt_result::set_matmul_desc_attribute(
                matmul_desc,
                cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSA,
                &transa as *const i32 as *const std::ffi::c_void,
                std::mem::size_of::<i32>(),
            ).map_err(|e| format!("set TRANSA: {e:?}"))?;
            cublaslt_result::set_matmul_desc_attribute(
                matmul_desc,
                cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSB,
                &transb as *const i32 as *const std::ffi::c_void,
                std::mem::size_of::<i32>(),
            ).map_err(|e| format!("set TRANSB: {e:?}"))?;

            // Matrix layouts
            let a_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, m as u64, k as i64)
                .map_err(|e| format!("A layout: {e:?}"))?;
            let b_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, k as i64)
                .map_err(|e| format!("B layout: {e:?}"))?;
            let c_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, m as i64)
                .map_err(|e| format!("C layout: {e:?}"))?;
            let d_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, m as i64)
                .map_err(|e| format!("D layout: {e:?}"))?;

            // Heuristic
            let matmul_pref = cublaslt_result::create_matmul_pref()
                .map_err(|e| format!("create_matmul_pref: {e:?}"))?;
            cublaslt_result::set_matmul_pref_attribute(
                matmul_pref,
                cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
                &ws_size as *const usize as *const std::ffi::c_void,
                std::mem::size_of::<usize>(),
            ).map_err(|e| format!("set pref ws: {e:?}"))?;

            let heuristic = cublaslt_result::get_matmul_algo_heuristic(
                lt_handle, matmul_desc, a_layout, b_layout, c_layout, d_layout, matmul_pref,
            ).map_err(|e| format!("heuristic: {e:?}"))?;

            eprintln!("    heuristic OK (ws_needed={})", heuristic.workspaceSize);

            // Execute matmul
            let cu_stream = stream as cublaslt_sys::cudaStream_t;
            let matmul_status = cublaslt_sys::cublasLtMatmul(
                lt_handle,
                matmul_desc,
                &alpha as *const f32 as *const std::ffi::c_void,
                d_w as *const std::ffi::c_void,
                a_layout,
                d_a as *const std::ffi::c_void,
                b_layout,
                &beta as *const f32 as *const std::ffi::c_void,
                d_c as *const std::ffi::c_void,
                c_layout,
                d_c as *mut std::ffi::c_void,
                d_layout,
                &heuristic.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t,
                ws_ptr as *mut std::ffi::c_void,
                ws_size,
                cu_stream,
            );

            // Sync
            cudarc::driver::sys::cuStreamSynchronize(stream);

            // Cleanup
            let _ = cublaslt_result::destroy_matmul_pref(matmul_pref);
            let _ = cublaslt_result::destroy_matrix_layout(d_layout);
            let _ = cublaslt_result::destroy_matrix_layout(c_layout);
            let _ = cublaslt_result::destroy_matrix_layout(b_layout);
            let _ = cublaslt_result::destroy_matrix_layout(a_layout);
            let _ = cublaslt_result::destroy_matmul_desc(matmul_desc);

            if matmul_status == cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
                Ok(())
            } else {
                Err(format!("cublasLtMatmul status={matmul_status:?}"))
            }
        };

        // Free buffers
        unsafe {
            cudarc::driver::sys::cuMemFree_v2(d_w);
            cudarc::driver::sys::cuMemFree_v2(d_a);
            cudarc::driver::sys::cuMemFree_v2(d_c);
        }

        result
    }

    #[test]
    fn test_cublaslt_all_dimensions() {
        let ctx = CudaContext::new(0).unwrap();
        let stream = ctx.default_stream();

        // Create cublasLt handle
        let lt_handle = cublaslt_result::create_handle().unwrap();

        // Allocate workspace via cuMemAlloc
        let ws_size: usize = 32 * 1024 * 1024;
        let mut ws_ptr: cudarc::driver::sys::CUdeviceptr = 0;
        unsafe {
            cudarc::driver::sys::cuMemAlloc_v2(&mut ws_ptr, ws_size);
        }

        let cases = [
            TestCase { m: 128, n: 64, k: 64, label: "train_small" },
            TestCase { m: 128, n: 512, k: 64, label: "eval_chunk" },
            TestCase { m: 128, n: 4096, k: 256, label: "h100_batch" },
            TestCase { m: 256, n: 4096, k: 256, label: "h100_h_s2" },
            TestCase { m: 128, n: 16384, k: 256, label: "h100_full" },
            TestCase { m: 51, n: 4096, k: 128, label: "h100_v_logits" },
            TestCase { m: 3, n: 5, k: 4, label: "tiny" },
        ];

        let cu_stream = stream.cu_stream();
        let mut failures = Vec::new();

        for tc in &cases {
            eprint!("  [{}] ({},{},{}) ... ", tc.label, tc.m, tc.n, tc.k);
            match test_cublaslt_matmul(lt_handle, cu_stream, ws_ptr, ws_size, tc) {
                Ok(()) => eprintln!("OK ✓"),
                Err(e) => {
                    eprintln!("FAIL ✗: {e}");
                    failures.push(format!("[{}] ({},{},{}): {e}", tc.label, tc.m, tc.n, tc.k));
                }
            }
        }

        // Cleanup
        unsafe {
            cudarc::driver::sys::cuMemFree_v2(ws_ptr);
            let _ = cublaslt_result::destroy_handle(lt_handle);
        }

        if !failures.is_empty() {
            panic!("cublasLtMatmul failures:\n{}", failures.join("\n"));
        }
    }
}
