from vllm.sequence import SequenceGroup, SequenceStatus
from vllm_mlu._mlu_utils import USE_PAGED
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.core.block_manager import SelfAttnBlockSpaceManager
from vllm.utils import Device
from vllm.logger import init_logger

logger = init_logger(__name__)


def vllm__core__block_manager__SelfAttnBlockSpaceManager__can_append_slots(self, seq_group: SequenceGroup,
                        num_lookahead_slots: int) -> bool:
    """Determine if there is enough space in the GPU KV cache to continue
    generation of the specified sequence group.

    We use a worst-case heuristic: assume each touched block will require a
    new allocation (either via CoW or new block). We can append slots if the
    number of touched blocks is less than the number of free blocks.

    "Lookahead slots" are slots that are allocated in addition to the slots
    for known tokens. The contents of the lookahead slots are not defined.
    This is used by speculative decoding when speculating future tokens.
    """
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: optimize the allocation strategy for unpagged mode
    '''
    if not USE_PAGED:
        return True
    else:
        num_touched_blocks = 0
        for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
            block_table = self.block_tables[seq.seq_id]

            num_touched_blocks += (
                block_table.get_num_blocks_touched_by_append_slots(
                    token_ids=block_table.get_unseen_token_ids(
                        seq.get_token_ids()),
                    num_lookahead_slots=num_lookahead_slots,
                ))

        num_free_gpu_blocks = self.block_allocator.get_num_free_blocks(
            Device.GPU)
        return num_touched_blocks <= num_free_gpu_blocks
    '''
    ==================
    End of MLU Hijack
    ==================
    '''


MluHijackObject.apply_hijack(SelfAttnBlockSpaceManager,
                             SelfAttnBlockSpaceManager.can_append_slots,
                             vllm__core__block_manager__SelfAttnBlockSpaceManager__can_append_slots)