from typing import Optional, Type
import torch
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.attention.backends.abstract import AttentionType
from vllm.attention.backends.mlu_attn import MLUFlashAttentionMetadata
from vllm_mlu.attention.backends.mlu_attn import MLUFlashAttentionImpl_V2

from .ring_attn import zigzag_ring_attn
from examples.cambricon_custom_func.vllm.mlu_hijack.distributed.parallel_state import (
    get_context_model_parallel_world_size)


vllm__attention__backends__flash_attn__MLUFlashAttentionImpl__forward_org = MLUFlashAttentionImpl_V2.forward

def vllm__attention__backends__flash_attn__MLUFlashAttentionImpl__forward_wraper(
    self,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    kv_cache: torch.Tensor,
    attn_metadata: MLUFlashAttentionMetadata,
    k_scale: float = 1.0,
    v_scale: float = 1.0,
    attn_type: AttentionType = AttentionType.DECODER,
    use_mla: bool = False,
) -> torch.Tensor:
    '''
    ==========================
    Modify by Context Parallel
    ==========================
    @brief: use ring attn when context parallel
    '''
    if get_context_model_parallel_world_size() > 1 and attn_metadata.prefill_metadata:
        return zigzag_ring_attn(self,
                                query=query.view(-1, self.num_heads, self.head_size),
                                key=key.view(-1, self.num_kv_heads, self.head_size),
                                value=value.view(-1, self.num_kv_heads, self.head_size),
                                kv_cache=kv_cache,
                                attn_metadata=attn_metadata)
    '''
    =======================
    End of Context Parallel
    =======================
    '''
    return vllm__attention__backends__flash_attn__MLUFlashAttentionImpl__forward_org(self,
                                                                                     query=query,
                                                                                     key=key,
                                                                                     value=value,
                                                                                     kv_cache=kv_cache,
                                                                                     attn_metadata=attn_metadata,
                                                                                     k_scale=k_scale,
                                                                                     v_scale=v_scale,
                                                                                     attn_type=attn_type)


MluHijackObject.apply_hijack(MLUFlashAttentionImpl_V2,
                             MLUFlashAttentionImpl_V2.forward,
                             vllm__attention__backends__flash_attn__MLUFlashAttentionImpl__forward_wraper)
