// 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 thiserror::Error;

use xgpu_common::{api_name::ApiFuncName, ipc::message::Argument};

use super::handler::*;

#[derive(Debug, Error)]
pub enum ServerErr {
    #[error("API run failed, api: {0}, error code: {1}")]
    ApiRunError(String, i32),

    #[error("Invalid argument type: {0}")]
    InvalidType(String),
}

pub trait ApiHandler: Send + Sync {
    fn handle_api(&self, args: &mut [Argument<'_>]) -> Result<Argument<'static>, ServerErr>;
}

macro_rules! dispatch_api {
    (
        $id:expr,
        { $($enum_variant:path => $handler:ident),* $(,)? }
    ) => {
        match $id as i32 {
            $(
                $enum_variant => Ok(&$handler as &'static dyn ApiHandler),
            )*
            _ => Err(ServerErr::ApiRunError(
                "Unknown API method ID".to_string(),
                $id as i32,
            )),
        }
    };
}

pub struct Api;

impl Api {
    #[inline]
    fn get_handler(id: u64) -> Result<&'static dyn ApiHandler, ServerErr> {
        dispatch_api!(id, {
            ApiFuncName::CUDADEVICERESET => CudaDeviceResetHandler,
            ApiFuncName::CUDADEVICESYNCHRONIZE => CudaDeviceSynchronizeHandler,
            ApiFuncName::CUDADEVICEGETSTREAMPRIORITYRANGE => CudaDeviceGetStreamPriorityRangeHandler,
            ApiFuncName::CUDAGETLASTERROR => CudaGetLastErrorHandler,
            ApiFuncName::CUDAPEEKATLASTERROR => CudaPeekAtLastErrorHandler,
            ApiFuncName::CUDAGETDEVICECOUNT => CudaGetDeviceCountHandler,
            ApiFuncName::CUDAGETDEVICEPROPERTIES_V2 => CudaGetDevicePropertiesV2Handler,
            ApiFuncName::CUDADEVICEGETATTRIBUTE => CudaDeviceGetAttributeHandler,
            ApiFuncName::CUDASETDEVICE => CudaSetDeviceHandler,
            ApiFuncName::CUDAGETDEVICE => CudaGetDeviceHandler,
            ApiFuncName::CUDASTREAMCREATEWITHPRIORITY => CudaStreamCreateWithPriorityHandler,
            ApiFuncName::CUDATHREADEXCHANGESTREAMCAPTUREMODE => CudaThreadExchangeStreamCaptureModeHandler,
            ApiFuncName::CUDASTREAMISCAPTURING => CudaStreamIsCapturingHandler,
            ApiFuncName::CUDASTREAMWAITEVENT => CudaStreamWaitEventHandler,
            ApiFuncName::CUDASTREAMSYNCHRONIZE => CudaStreamSynchronizeHandler,
            ApiFuncName::NCCLGETVERSION => NcclGetVersionHandler,
            ApiFuncName::CUDAEVENTCREATEWITHFLAGS => CudaEventCreateWithFlagsHandler,
            ApiFuncName::CUDAEVENTRECORD => CudaEventRecordHandler,
            ApiFuncName::CUDAEVENTRECORDWITHFLAGS => CudaEventRecordWithFlagsHandler,
            ApiFuncName::CUDAEVENTQUERY => CudaEventQueryHandler,
            ApiFuncName::CUDAEVENTSYNCHRONIZE => CudaEventSynchronizeHandler,
            ApiFuncName::CUDAEVENTDESTROY => CudaEventDestroyHandler,
            ApiFuncName::CUDAEVENTELAPSEDTIME => CudaEventElapsedTimeHandler,
            ApiFuncName::CUDAMEMSET => CudaMemsetHandler,
            ApiFuncName::CUDAPOINTERGETATTRIBUTES => CudaPointerGetAttributesHandler,
            ApiFuncName::CUDEVICEGET => CuDeviceGetHandler,
            ApiFuncName::CUDEVICEGETATTRIBUTE => CuDeviceGetAttributeHandler,
            ApiFuncName::CUDEVICEPRIMARYCTXGETSTATE => CuDevicePrimaryCtxGetStateHandler,
            ApiFuncName::CUCTXGETCURRENT => CuCtxGetCurrentHandler,
            ApiFuncName::CUFUNCGETATTRIBUTE => CuFuncGetAttributeHandler,
            ApiFuncName::CUDAMEMSETASYNC => CudaMemsetAsyncHandler,
            ApiFuncName::CUDAFUNCSETATTRIBUTE => CudaFuncSetAttributeHandler,
            ApiFuncName::CUDAFUNCGETATTRIBUTES => CudaFuncGetAttributesHandler,
            ApiFuncName::CUBLASCREATE_V2 => CublasCreateV2handler,
            ApiFuncName::CUBLASSETWORKSPACE_V2 => CublasSetWorkspaceV2handler,
            ApiFuncName::CUBLASSETSTREAM_V2 => CublasSetStreamV2handler,
            ApiFuncName::CUBLASGETSTREAM_V2 => CublasGetStreamV2handler,
            ApiFuncName::CUBLASLTMATMULPREFERENCECREATE => CublasLtMatmulPreferenceCreateHandler,
            ApiFuncName::CUBLASLTMATMULDESCCREATE => CublasLtMatmulDescCreateHandler,
            ApiFuncName::CUBLASLTMATMULDESCSETATTRIBUTE => CublasLtMatmulDescSetAttributeHandler,
            ApiFuncName::CUBLASLTMATRIXLAYOUTCREATE => CublasLtMatrixLayoutCreateHandler,
            ApiFuncName::CUBLASLTMATMULPREFERENCESETATTRIBUTE => CublasLtMatmulPreferenceSetAttributeHandler,
            ApiFuncName::CUBLASGETMATHMODE => CublasGetMathModeHandler,
            ApiFuncName::CUBLASSETMATHMODE => CublasSetMathModeHandler,
            ApiFuncName::CUBLASLTMATMULALGOGETHEURISTIC => CublasLtMatmulAlgoGetHeuristicHandler,
            ApiFuncName::CUBLASLTMATMUL => CublasLtMatmulHandler,
            ApiFuncName::CUBLASLTMATRIXLAYOUTDESTROY => CublasLtMatrixLayoutDestroyHandler,
            ApiFuncName::CUBLASLTMATMULDESCDESTROY => CublasLtMatmulDescDestroyHandler,
            ApiFuncName::CUBLASLTMATMULPREFERENCEDESTROY => CublasLtMatmulPreferenceDestroyHandler,
            ApiFuncName::CUBLASSGEMM_V2 => CublasSgemmV2handler,
            ApiFuncName::CUDAMEMCPYASYNC => CudaMemcpyAsyncHandler,
            ApiFuncName::CUDAMALLOC => CudaMallocHandler,
            ApiFuncName::CUDAFREE => CudaFreeHandler,
            ApiFuncName::NVMLINIT_V2 => NvmlInitV2Handler,
            ApiFuncName::NVMLSHUTDOWN => NvmlShutdownHandler,
            ApiFuncName::NVMLDEVICEGETCOUNT_V2 => NvmlDeviceGetCountV2handler,
            ApiFuncName::NCCLCOMMDESTROY => NcclCommDestroyHandler,
            ApiFuncName::NCCLCOMMABORT => NcclCommAbortHandler,
            ApiFuncName::NCCLCOMMFINALIZE => NcclCommFinalizeHandler,
            ApiFuncName::NCCLCOMMSPLIT => NcclCommSplitHandler,
            ApiFuncName::NCCLCOMMGETASYNCERROR => NcclCommGetAsyncErrorHandler,
            ApiFuncName::NCCLBCAST => NcclBcastHandler,
            ApiFuncName::NCCLALLREDUCE => NcclAllReduceHandler,
            ApiFuncName::NCCLALLGATHER => NcclAllGatherHandler,
            ApiFuncName::NCCLSEND => NcclSendHandler,
            ApiFuncName::NCCLRECV => NcclRecvHandler,
            ApiFuncName::NCCLGROUPSTART => NcclGroupStartHandler,
            ApiFuncName::NCCLGROUPEND => NcclGroupEndHandler,
            ApiFuncName::NCCLGETUNIQUEID => NcclGetUniqueIdHandler,
            ApiFuncName::NCCLCOMMINITRANKCONFIG => NcclCommInitRankConfigHandler,
            ApiFuncName::CUDALAUNCHKERNEL => CudaLaunchKernelHandler,
            ApiFuncName::NCCLREDUCE => NcclReduceHandler,
            ApiFuncName::CUBLASGEMMEX => CublasGemmExHandler,
            ApiFuncName::CUBLASSGEMMSTRIDEDBATCHED => CublasSgemmStridedBatchedHandler,
            ApiFuncName::CUBLASGEMMSTRIDEDBATCHEDEX => CublasGemmStridedBatchedExHandler,
            ApiFuncName::CUBLASLTCREATE => CublasLtCreateHandler,
            ApiFuncName::CUBLASLTDESTROY => CublasLtDestroyHandler,
            ApiFuncName::CUDAOCCUPANCYMAXACTIVEBLOCKSPERMULTIPROCESSORWITHFLAGS => CudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlagsHandler,
            ApiFuncName::LOADDYNLIBRARY => LoadLibraryHandler,
            ApiFuncName::NCCLCOMMSHRINK => NcclCommShrinkHandler,
            ApiFuncName::NCCLCOMMINITNEWRANK => NcclCommInitNewRankHandler,
            ApiFuncName::NCCLCOMMADDNEWRANK => NcclCommAddNewRankHandler,
            ApiFuncName::NCCLCOMMSETUPNEWRANK => NcclCommSetupNewRankHandler,
            ApiFuncName::CUDAHOSTALLOC => CudaHostAllocHandler,
        })
    }

    pub fn invoke(id: u64, args: &mut [Argument<'_>]) -> Result<Argument<'static>, ServerErr> {
        Self::get_handler(id)?.handle_api(args)
    }
}