// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use xgpu_macros::api_hook;

#[api_hook(CublasApi, backend = crate::hook_impl::ipc::cublas::CublasApiImpl)]
mod api {
    use std::ffi::{c_int, c_longlong, c_void};

    use cudax::cublas::*;

    unsafe extern "C" {
        pub fn cublasCreate_v2(handle: *mut cublasHandle_t) -> cublasStatus_t;

        pub fn cublasSetWorkspace_v2(
            handle: cublasHandle_t,
            workspace: *mut c_void,
            workspace_size_in_bytes: usize,
        ) -> cublasStatus_t;

        pub fn cublasSetStream_v2(
            handle: cublasHandle_t,
            stream_id: cudaStream_t,
        ) -> cublasStatus_t;

        pub fn cublasGetStream_v2(
            handle: cublasHandle_t,
            stream_id: *mut cudaStream_t,
        ) -> cublasStatus_t;

        pub fn cublasGetMathMode(handle: cublasHandle_t, mode: *mut cublasMath_t)
            -> cublasStatus_t;

        pub fn cublasSetMathMode(handle: cublasHandle_t, mode: cublasMath_t) -> cublasStatus_t;

        pub fn cublasSgemm_v2(
            handle: cublasHandle_t,
            transa: cublasOperation_t,
            transb: cublasOperation_t,
            m: c_int,
            n: c_int,
            k: c_int,
            alpha: *const f32,
            a: *const f32,
            lda: c_int,
            b: *const f32,
            ldb: c_int,
            beta: *const f32,
            c: *mut f32,
            ldc: c_int,
        ) -> cublasStatus_t;

        pub fn cublasGemmEx(
            handle: cublasHandle_t,
            transa: cublasOperation_t,
            transb: cublasOperation_t,
            m: c_int,
            n: c_int,
            k: c_int,
            alpha: *const c_void,
            a: *const c_void,
            atype: cudaDataType,
            lda: c_int,
            b: *const c_void,
            btype: cudaDataType,
            ldb: c_int,
            beta: *const c_void,
            c: *mut c_void,
            ctype: cudaDataType,
            ldc: c_int,
            compute_type: cublasComputeType_t,
            algo: cublasGemmAlgo_t,
        ) -> cublasStatus_t;

        pub fn cublasSgemmStridedBatched(
            handle: cublasHandle_t,
            transa: cublasOperation_t,
            transb: cublasOperation_t,
            m: c_int,
            n: c_int,
            k: c_int,
            alpha: *const f32,
            a: *const f32,
            lda: c_int,
            stride_a: c_longlong,
            b: *const f32,
            ldb: c_int,
            stride_b: c_longlong,
            beta: *const f32,
            c: *mut f32,
            ldc: c_int,
            stride_c: c_longlong,
            batch_count: c_int,
        ) -> cublasStatus_t;

        pub fn cublasGemmStridedBatchedEx(
            handle: cublasHandle_t,
            transa: cublasOperation_t,
            transb: cublasOperation_t,
            m: c_int,
            n: c_int,
            k: c_int,
            alpha: *const c_void,
            a: *const c_void,
            atype: cudaDataType,
            lda: c_int,
            stride_a: c_longlong,
            b: *const c_void,
            btype: cudaDataType,
            ldb: c_int,
            stride_b: c_longlong,
            beta: *const c_void,
            c: *mut c_void,
            ctype: cudaDataType,
            ldc: c_int,
            stride_c: c_longlong,
            batch_count: c_int,
            compute_type: cublasComputeType_t,
            algo: cublasGemmAlgo_t,
        ) -> cublasStatus_t;
    }
}

pub use api::CublasApi;