# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/_custom_ops.py
import logging
from typing import List, Optional, Tuple

import torch

from sglang.srt.utils import get_bool_env_var, is_hip, is_hpu, is_npu
try:
    from lmslim import quant_ops 
    from lmslim import quant_tools 
except Exception:
    print("INFO: Please install lmslim if you want to infer gptq or awq  or w8a8 model.\n") 
try:
    import lightop
except Exception:
    print("INFO: Please install lightop if you want to infer awq of marlin.\n") 

logger = logging.getLogger(__name__)
use_vllm_custom_allreduce = get_bool_env_var(
    "USE_VLLM_CUSTOM_ALLREDUCE", default="false"
)
use_dcu_custom_allreduce = get_bool_env_var(
    "USE_DCU_CUSTOM_ALLREDUCE", default="true"
)

if not is_hpu():
    # ROCm does not use vllm custom allreduce
    if use_vllm_custom_allreduce and not is_hip():
        try:
            import vllm._C  # noqa: F401
        except ImportError as e:
            logger.warning("Failed to import from vllm._C with %r", e)
    else:
        try:
            import sgl_kernel
        except ImportError as e:
            logger.warning("Failed to import from custom_ar with %r", e)

if use_dcu_custom_allreduce:
    try:
        import vllm._C
    except ImportError as e:
        logger.warning("Failed to import from vllm._C with %r", e)

if not is_hip() and not is_npu():
    if use_vllm_custom_allreduce:
        custom_op = torch.ops._C_custom_ar
    else:
        custom_op = sgl_kernel.allreduce

    # custom allreduce
    def init_custom_ar(
        ipc_tensors: List[torch.Tensor],
        rank_data: torch.Tensor,
        rank: int,
        full_nvlink: bool,
    ) -> int:
        return custom_op.init_custom_ar(ipc_tensors, rank_data, rank, full_nvlink)

    def all_reduce(
        fa: int,
        inp: torch.Tensor,
        out: torch.Tensor,
        reg_buffer: int,
        reg_buffer_sz_bytes: int,
    ) -> None:
        custom_op.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes)

    def dispose(fa: int) -> None:
        custom_op.dispose(fa)

    def meta_size() -> int:
        return custom_op.meta_size()

    def register_buffer(fa: int, ipc_tensors: List[int]) -> None:
        return custom_op.register_buffer(fa, ipc_tensors)

    def get_graph_buffer_ipc_meta(fa: int) -> Tuple[List[int], List[int]]:
        return custom_op.get_graph_buffer_ipc_meta(fa)

    def register_graph_buffers(
        fa: int, handles: List[List[int]], offsets: List[List[int]]
    ) -> None:
        custom_op.register_graph_buffers(fa, handles, offsets)

elif is_hip and use_dcu_custom_allreduce:
    # custom ar
    def init_custom_ar(ipc_tensors: list[torch.Tensor], rank_data: torch.Tensor,
                    rank: int, fully_connected: bool) -> int:
        return torch.ops._C_custom_ar.init_custom_ar(ipc_tensors, rank_data, rank,
                                                    fully_connected)


    def all_reduce(fa: int, inp: torch.Tensor, out: torch.Tensor, reg_buffer: int,
                reg_buffer_sz_bytes: int) -> None:
        torch.ops._C_custom_ar.all_reduce(fa, inp, out, reg_buffer,
                                        reg_buffer_sz_bytes)


    def dispose(fa: int) -> None:
        torch.ops._C_custom_ar.dispose(fa)


    def meta_size() -> int:
        return torch.ops._C_custom_ar.meta_size()


    def register_buffer(fa: int, ipc_tensors: list[int]) -> None:
        return torch.ops._C_custom_ar.register_buffer(fa, ipc_tensors)


    def get_graph_buffer_ipc_meta(fa: int) -> tuple[list[int], list[int]]:
        return torch.ops._C_custom_ar.get_graph_buffer_ipc_meta(fa)


    def register_graph_buffers(fa: int, handles: list[list[int]],
                            offsets: list[list[int]]) -> None:
        torch.ops._C_custom_ar.register_graph_buffers(fa, handles, offsets)


    def allocate_shared_buffer_and_handle(size: int) -> tuple[int, torch.Tensor]:
        return torch.ops._C_custom_ar.allocate_shared_buffer_and_handle(size)


    def open_mem_handle(mem_handle: torch.Tensor):
        return torch.ops._C_custom_ar.open_mem_handle(mem_handle)


    def free_shared_buffer(ptr: int) -> None:
        torch.ops._C_custom_ar.free_shared_buffer(ptr)


    def read_cache(
            keys: torch.Tensor,
            values: torch.Tensor,
            key_caches: list[torch.Tensor],
            value_caches: list[torch.Tensor],
            slot_mapping: torch.Tensor,
            kv_cache_dtype: str
    ) -> None:
        torch.ops._C_cache_ops.read_cache(keys, values, key_caches,
                                        value_caches, slot_mapping,
                                        kv_cache_dtype)

    def write_cache_multi_layers(
            keys: torch.Tensor,
            values: torch.Tensor,
            key_caches: list[torch.Tensor],
            value_caches: list[torch.Tensor],
            slot_mapping: torch.Tensor,
            kv_cache_dtype: str
    ) -> None:
        torch.ops._C_cache_ops.write_cache_multi_layers(keys, values, key_caches,
                                                        value_caches, slot_mapping,
                                                        kv_cache_dtype)

else:
    # sgl_kernel ROCM custom allreduce

    def init_custom_ar(
        meta: torch.Tensor,
        rank_data: torch.Tensor,
        handles: List[str],
        offsets: List[int],
        rank: int,
        full_nvlink: bool,
    ) -> int:
        return sgl_kernel.allreduce.init_custom_ar(
            meta, rank_data, handles, offsets, rank, full_nvlink
        )

    def all_reduce_reg(fa: int, inp: torch.Tensor, out: torch.Tensor) -> None:
        sgl_kernel.allreduce.all_reduce_reg(fa, inp, out)

    def all_reduce_unreg(
        fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor
    ) -> None:
        sgl_kernel.allreduce.all_reduce_unreg(fa, inp, reg_buffer, out)

    def dispose(fa: int) -> None:
        sgl_kernel.allreduce.dispose(fa)

    def meta_size() -> int:
        return sgl_kernel.allreduce.meta_size()

    def register_buffer(
        fa: int, t: torch.Tensor, handles: List[str], offsets: List[int]
    ) -> None:
        return sgl_kernel.allreduce.register_buffer(fa, t, handles, offsets)

    def get_graph_buffer_ipc_meta(fa: int) -> Tuple[torch.Tensor, List[int]]:
        return sgl_kernel.allreduce.get_graph_buffer_ipc_meta(fa)

    def register_graph_buffers(
        fa: int, handles: List[str], offsets: List[List[int]]
    ) -> None:
        sgl_kernel.allreduce.register_graph_buffers(fa, handles, offsets)

    def allocate_meta_buffer(size: int) -> torch.Tensor:
        return sgl_kernel.allreduce.allocate_meta_buffer(size)

    def get_meta_buffer_ipc_handle(inp: torch.Tensor) -> torch.Tensor:
        return sgl_kernel.allreduce.get_meta_buffer_ipc_handle(inp)

    # ROCM custom quick allreduce

    def init_custom_qr(
        rank: int, world_size: int, qr_max_size: Optional[int] = None
    ) -> int:
        return sgl_kernel.allreduce.init_custom_qr(world_size, rank, qr_max_size)

    def qr_get_handle(fa: int) -> torch.Tensor:
        return sgl_kernel.allreduce.qr_get_handle(fa)

    def qr_open_handles(fa: int, handles: list[torch.Tensor]) -> None:
        sgl_kernel.allreduce.qr_open_handles(fa, handles)

    def qr_all_reduce(
        fa: int,
        inp: torch.Tensor,
        out: torch.Tensor,
        quant_level: int,
        cast_bf2half: bool,
    ) -> None:
        sgl_kernel.allreduce.qr_all_reduce(fa, inp, out, quant_level, cast_bf2half)

    def qr_destroy(fa: int) -> None:
        sgl_kernel.allreduce.qr_destroy(fa)

    def qr_max_size() -> int:
        return sgl_kernel.allreduce.qr_max_size()


def mscclpp_generate_unique_id() -> bytes:
    return sgl_kernel.allreduce.mscclpp_generate_unique_id()


def mscclpp_init_context(
    unique_id: bytes,
    rank: int,
    world_size: int,
    scratch: torch.Tensor,
    put_buffer: torch.Tensor,
    nranks_per_node: int,
    rank_to_node: List[int],
    rank_to_ib: List[int],
    context_selection: int,
) -> int:
    return sgl_kernel.allreduce.mscclpp_init_context(
        unique_id,
        rank,
        world_size,
        scratch,
        put_buffer,
        nranks_per_node,
        rank_to_node,
        rank_to_ib,
        context_selection,
    )


def mscclpp_allreduce(
    context: int, inp: torch.Tensor, out: torch.Tensor, nthreads: int, nblocks: int
) -> None:
    return sgl_kernel.allreduce.mscclpp_allreduce(context, inp, out, nthreads, nblocks)

def triton_scaled_mm(a: torch.Tensor,
                      b: torch.Tensor,
                      scale_a: torch.Tensor,
                      scale_b: torch.Tensor,
                      out_dtype: torch.dtype,
                      bias: Optional[torch.Tensor] = None,
                      best_config:Optional[list] = None) -> torch.Tensor:

    return quant_ops.triton_scaled_mm(a, b,scale_a,scale_b,out_dtype,bias,best_config)

def cutlass_scaled_mm(a: torch.Tensor,
                      b: torch.Tensor,
                      scale_a: torch.Tensor,
                      scale_b: torch.Tensor,
                      out_dtype: torch.dtype,
                      bias: Optional[torch.Tensor] = None) -> torch.Tensor:
    """
    `cutlass_scaled_mm` implements a fused version of
        `output = torch.mm((scale_a * a), (scale_b * b)).to(out_dtype)`
    where scale_a * a and scale_b * b are implemented using numpy-style
    broadcasting.

    In order to support blockwise scaling like found in DeepSeek V3 we also
    support extended "group" broadcast rules. We extend the numpy-style
    broadcasting rules with the following rule:
        "if the extent of a dimension in the source shape is between 1 and
        corresponding extent in the target shape we repeat each element along
        that dimension  src_shape[dim] // target_shape[dim] times consecutively"
    example if we have:
          a = [[1, 2], and target_shape = (2, 4)
               [3, 4]]
    then we would expand a to:
          a = [[1, 1, 2, 2],
               [3, 3, 4, 4]]
    currently we only support the case:
        scale_a.shape * [1, 128] == a.shape
        scale_b.shape * [128, 128] == b.shape
    """
    assert (out_dtype is torch.bfloat16 or out_dtype is torch.float16)
    assert bias is None or bias.shape[0] == b.shape[
        1] and bias.dtype == out_dtype

    # m = a.shape[0]
    # n = b.shape[1]

    # cutlass_compatible_b = (b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0)
    # if current_platform.is_rocm() or not cutlass_compatible_b:
    #     from vllm.model_executor.layers.quantization.compressed_tensors.triton_scaled_mm import (  # noqa
    #         triton_scaled_mm)
    #     return triton_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)

    # out = torch.empty((m, n), dtype=out_dtype, device=a.device)

    # torch.ops._C.cutlass_scaled_mm(out, a, b, scale_a, scale_b, bias)

    # return out
    #return quant_ops.cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
    return quant_ops.rocblas_scaled_mm_nn(a, b, scale_a, scale_b, out_dtype, bias)

def rocblas_scaled_mm(a: torch.Tensor,
                      b: torch.Tensor,
                      scale_a: torch.Tensor,
                      scale_b: torch.Tensor,
                      out_dtype: torch.dtype,
                      bias: Optional[torch.Tensor] = None) -> torch.Tensor:

    return quant_ops.rocblas_scaled_mm_nn(a, b, scale_a, scale_b, out_dtype, bias)

def blaslt_scaled_mm(a: torch.Tensor,
                      b: torch.Tensor,
                      scale_a: torch.Tensor,
                      scale_b: torch.Tensor,
                      out_dtype: torch.dtype,
                      bias: Optional[torch.Tensor] = None) -> torch.Tensor:
    m = a.shape[0]
    n = b.shape[0]
    k = a.shape[1]
    _, out = quant_ops.hipblaslt_w8a8_gemm(a, b, scale_a, scale_b, m, n, k, 'NT', out_dtype)
    return out

def triton_int8_gemm_helper(m: int,
                             n: int,
                             k: int,
                             per_token_act_quant: bool,
                             per_out_channel_weight_quant: bool,
                             use_bias: bool,
                             out_dtype: type[torch.dtype] = torch.float16,
                             device: str = "cuda:0",
                             best_config:Optional[list] = None,
                             repeat:Optional[int] = 2):
    return quant_tools.triton_int8_gemm_helper(m,n,k,per_token_act_quant,per_out_channel_weight_quant,use_bias,out_dtype,device,best_config,repeat)