from typing import Union

import torch

from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.model_executor.layers.spec_decode_base_sampler import SpecDecodeBaseSampler
from vllm.platforms import current_platform


def vllm__model_executor__layers__spec_decode_base_sampler__SpecDecodeBaseSampler__init_gpu_tensors(
    self, device: Union[int, str]
) -> None:
    assert self.num_accepted_tokens is None
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: Add mlu device support.
    '''
    if isinstance(device, int) and current_platform.is_mlu():
        device = f"mlu:{device}"
    elif isinstance(device, int) and current_platform.is_cuda():
        device = f"cuda:{device}"
    elif not isinstance(device, str):
        raise ValueError(f"Device must be int or str, get {type(device)}")
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    self.num_accepted_tokens = torch.tensor(0,
                                            dtype=torch.long,
                                            device=device)
    self.num_emitted_tokens = torch.tensor(0,
                                           dtype=torch.long,
                                           device=device)

MluHijackObject.apply_hijack(
    SpecDecodeBaseSampler,
    SpecDecodeBaseSampler.init_gpu_tensors,
    vllm__model_executor__layers__spec_decode_base_sampler__SpecDecodeBaseSampler__init_gpu_tensors
)
