import torch
from typing import List, Optional
from vllm import _mlu_ops as mlu_ops
from vllm_mlu._mlu_utils import *
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.distributed import get_tensor_model_parallel_rank
from vllm.distributed.parallel_state import get_tensor_model_parallel_group
from vllm.model_executor.models.mixtral_quant import MixtralAttention
from vllm_mlu.model_executor.layers.quantization.smoothquant import SmoothQuantLinearMethod
from vllm.attention.backends.abstract import (AttentionMetadata,
                                              AttentionType)
from vllm.attention.backends.utils import get_num_prefill_decode_query_kv_tokens
from vllm.forward_context import get_forward_context
from vllm.utils import direct_register_custom_op
from vllm.attention.backends.mlu_attn import (MLUFlashAttentionMetadata,  
                                              _get_query_key_seq_metadata, 
                                              _get_causal_option)


def vllm__model_executor__models__mixtral__MixtralAttention__forward(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
    kv_cache: torch.Tensor,
    attn_metadata: AttentionMetadata,
    residual: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    qkv, _ = self.qkv_proj(hidden_states)
    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: pack q & k to fit tmo.apply_rotary
    ''' 
    qk, _ = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
    self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
    '''
    ==================
    End of MLU Hijack
    ==================
    '''

    '''
    =====================================================
    Modify by Context Communication Computation Parallel
    =====================================================
    @brief: call flash_attn_sq_mm_allreduce to finish forward
    '''
    if (attn_metadata.prefill_metadata) and \
       (kv_cache[0].numel() > 0) and \
       (hasattr(self.o_proj, 'quant_method')) and \
       (isinstance(self.o_proj.quant_method, SmoothQuantLinearMethod)) and \
       (self.o_proj.quant_method.quant_config.input_quant_method == "per_token"):
        rank = get_tensor_model_parallel_rank()
        pg = get_tensor_model_parallel_group().device_group
        cncl_comm = pg._get_backend(torch.device("mlu")).get_cncl_comm(rank)

        return torch.ops.vllm.context_attn_comm_cmpt_parallel_flash_attention_v2(
                    q, k, v,
                    self.num_heads, self.head_dim, self.num_kv_heads,
                    kv_cache, self.attn.impl.kv_cache_dtype,
                    1.0, 1.0, self.scaling,
                    cncl_comm,
                    self.o_proj.smooth, self.o_proj.qweight, 
                    self.o_proj.per_channel_scale.to(torch.float),
                    self.o_proj.quant_method.parallel_num,
                    residual, self.attn.impl.sliding_window, self.attn.impl.alibi_slopes
                )
    '''
    =====================================================
    End of Context Communication Computation Parallel
    =====================================================
    '''

    attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: add residual
    ''' 
    output, _ = self.o_proj(attn_output, residual)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    return output


MluHijackObject.apply_hijack(MixtralAttention,
                             MixtralAttention.forward,
                             vllm__model_executor__models__mixtral__MixtralAttention__forward)


def context_attn_comm_cmpt_parallel_flash_attention_v2(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    num_heads: int,
    head_size: int,
    num_kv_heads: int,
    kv_cache: List[torch.Tensor],
    kv_cache_dtype: str,
    k_scale: float,
    v_scale: float,
    softmax_scale: float,
    cncl_comm: int,
    smooth: torch.Tensor,
    qweight: torch.Tensor,
    per_channel_scale: torch.Tensor,
    parallel_num: int,
    residual: Optional[torch.Tensor] = None,
    window_size: Optional[List[int]] = None,
    alibi_slopes: Optional[torch.Tensor] = None,
    logits_soft_cap: Optional[float] = None,
) -> torch.Tensor:
    current_metadata = get_forward_context()
    assert current_metadata is not None
    assert isinstance(current_metadata, MLUFlashAttentionMetadata)
    attn_metadata: MLUFlashAttentionMetadata = current_metadata

    # Reshape the query, key, and value tensors.
    query = query.view(-1, num_heads, head_size)
    if (key is not None) and (key is not None):
        key = key.view(-1, num_kv_heads, head_size)
        value = value.view(-1, num_kv_heads, head_size)

    kv_cache_, kv_cache_scale_ = kv_cache
    key_cache = kv_cache_[0]
    value_cache = kv_cache_[1]
    key_cache_scale, value_cache_scale = None, None
    if kv_cache_scale_.numel() > 0:
        key_cache_scale = kv_cache_scale_[0]
        value_cache_scale = kv_cache_scale_[1]

    # if not specified in self.attn.forward params, use default DECODER
    attn_type = AttentionType.DECODER

    # We skip updating the KV cache under two conditions:
    #  a. When the Attention Type is ENCODER. In this phase, we compute
    #     only the encoder attention without updating the cache.
    #  b. When both Key and Value are None. This occurs during
    #     cross-attention computation in the decoding phase, where the KV
    #     cache is already populated with the cross-attention tensor.
    #     Thus, we skip cache updates during this time.
    if (attn_type != AttentionType.ENCODER) and (key is not None) and (
            value is not None):
        if attn_type == AttentionType.ENCODER_DECODER:
            # Update cross-attention KV cache (prefill-only)
            updated_slot_mapping = attn_metadata.cross_slot_mapping
        else:
            # Update self-attention KV cache (prefill/decode)
            updated_slot_mapping = attn_metadata.slot_mapping

    # Reshape the input keys and values and store them in the cache.
    # If kv_cache is not provided, the new key and value tensors are
    # not cached. This happens during the initial memory profiling run.
    if USE_PAGED:
        if kv_cache_dtype == 'int8':
            mlu_ops.quant_to_paged_cache(key,
                                        value,
                                        key_cache,
                                        value_cache,
                                        key_cache_scale,
                                        value_cache_scale,
                                        attn_metadata.slot_mapping.flatten())
        else:
            mlu_ops.reshape_paged_cache(key,
                                        value,
                                        key_cache,
                                        value_cache,
                                        attn_metadata.slot_mapping.flatten())
    else:
        # FIXME: After TMO-1496 is completed, remove this code.
        if key.stride() != value.stride():
            key = key.contiguous()
            value = value.contiguous()
        if kv_cache_dtype == 'int8':
            mlu_ops.quant_to_linear_cache(key,
                                            value,
                                            key_cache,
                                            value_cache,
                                            key_cache_scale,
                                            value_cache_scale,
                                            attn_metadata.cu_seq_lens,
                                            attn_metadata.max_seq_len,
                                            True, # packed
                                            None, # context_seq_offset
                                            attn_metadata.batch_ids,
                                            attn_metadata.slot_mapping_unpaged)
        else:
            mlu_ops.reshape_linear_cache(key,
                                            value,
                                            key_cache,
                                            value_cache,
                                            attn_metadata.cu_seq_lens,
                                            attn_metadata.max_seq_len,
                                            True, # packed
                                            None, # context_seq_offset
                                            attn_metadata.batch_ids,
                                            attn_metadata.slot_mapping_unpaged)

    (num_prefill_query_tokens, num_prefill_kv_tokens,
    num_decode_query_tokens) = \
        get_num_prefill_decode_query_kv_tokens(attn_metadata, attn_type)
    decode_query = query[num_prefill_query_tokens:]
    # QKV for prefill.
    query = query[:num_prefill_query_tokens]
    assert query.shape[0] == num_prefill_query_tokens
    assert decode_query.shape[0] == num_decode_query_tokens

    alibi_slopes = None if alibi_slopes is None else \
                            alibi_slopes.repeat(attn_metadata.num_prefills, 1)
    prefill_meta = attn_metadata.prefill_metadata
    # Prompt run.
    if (prefill_meta.block_tables is None
            or prefill_meta.block_tables.numel() == 0):
        # normal attention
        # When block_tables are not filled, it means q and k are the
        # prompt, and they have the same length.
        q_seq_start_loc, q_seq_len, k_seq_start_loc, k_seq_len = \
            _get_query_key_seq_metadata(prefill_meta, True, attn_type)

        key = key[:num_prefill_kv_tokens]
        value = value[:num_prefill_kv_tokens]

        output = mlu_ops.flash_attn_sq_mm_allreduce(cncl_comm, 
                                    query, key, value, 
                                    q_seq_start_loc, k_seq_start_loc, 
                                    alibi_slopes, None, 
                                    smooth, qweight, 
                                    per_channel_scale, None,
                                    q_seq_len, k_seq_len,
                                    softmax_scale, _get_causal_option(attn_type), 
                                    -1 if window_size is None \
                                        else window_size[0],
                                    -1 if window_size is None \
                                        else window_size[1],
                                    torch.float, parallel_num)
    else:
        # prefix-enabled attention
        assert attn_type == AttentionType.DECODER, (
            "Only decoder-only models support prefix caching")
        assert prefill_meta.seq_lens is not None
        max_seq_len = max(prefill_meta.seq_lens)
        output = mlu_ops.flash_attn_sq_mm_allreduce(cncl_comm, 
                                    query, key_cache, value_cache, 
                                    prefill_meta.query_start_loc, prefill_meta.seq_start_loc, 
                                    alibi_slopes, None, 
                                    smooth, qweight, 
                                    per_channel_scale, None,
                                    prefill_meta.max_query_len, max_seq_len,
                                    softmax_scale, True, 
                                    -1 if window_size is None \
                                        else window_size[0],
                                    -1 if window_size is None \
                                        else window_size[1],
                                    torch.float, parallel_num)

    # Add residual.
    if residual is not None:
        output = output + residual
    return output


def context_attn_comm_cmpt_parallel_flash_attention_v2_fake(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    num_heads: int,
    head_size: int,
    num_kv_heads: int,
    kv_cache: List[torch.Tensor],
    kv_cache_dtype: str,
    k_scale: float,
    v_scale: float,
    softmax_scale: float,
    cncl_comm: int,
    smooth: torch.Tensor,
    qweight: torch.Tensor,
    per_channel_scale: torch.Tensor,
    parallel_num: int,
    residual: Optional[torch.Tensor] = None,
    window_size: Optional[List[int]] = None,
    alibi_slopes: Optional[torch.Tensor] = None,
    logits_soft_cap: Optional[float] = None,
) -> torch.Tensor:
    return torch.empty_like(query)


direct_register_custom_op(
    op_name="context_attn_comm_cmpt_parallel_flash_attention_v2",
    op_func=context_attn_comm_cmpt_parallel_flash_attention_v2,
    mutates_args=["kv_cache"],
    fake_impl=context_attn_comm_cmpt_parallel_flash_attention_v2_fake,
)