import torch

from typing import Dict, List, Optional, Union, Any
from transformers import LlamaConfig
from vllm.attention import AttentionMetadata
from vllm.config import CacheConfig
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm_mlu.model_executor.layers.feed_forward import FeedForward
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.logger import init_logger

from vllm.model_executor.models.llama import LlamaAttention, LlamaDecoderLayer, LlamaModel
from vllm.sequence import IntermediateTensors

from vllm_mlu.model_executor.models.layer_utils import (
    decoder_layer_forward_base, decoder_model_forward_base_pp,
    is_per_tensor_smoothquant, is_per_token_smoothquant,
    quant_fusion_with_rmsnorm)

logger = init_logger(__name__)

vllm__module_executor__models__llama__LlamaAttention__init__org = LlamaAttention.__init__


def vllm__module_executor__models__llama__LlamaAttention____init__(
    self,
    config: LlamaConfig,
    hidden_size: int,
    num_heads: int,
    num_kv_heads: int,
    rope_theta: float = 10000,
    rope_scaling: Optional[Dict[str, Any]] = None,
    max_position_embeddings: int = 8192,
    quant_config: Optional[QuantizationConfig] = None,
    bias: bool = False,
    cache_config: Optional[CacheConfig] = None,
    prefix: str = "",
) -> None:
    vllm__module_executor__models__llama__LlamaAttention__init__org(
        self,
        config,
        hidden_size,
        num_heads,
        num_kv_heads,
        rope_theta,
        rope_scaling,
        max_position_embeddings,
        quant_config,
        bias,
        cache_config,
        prefix)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: add rope_scaling params
    ''' 
    self.rope_scaling = rope_scaling
    '''
    ==================
    End of MLU Hijack
    ==================
    '''


def vllm__module_executor__models__llama__LlamaAttention__forward(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
    kv_cache: torch.Tensor,
    attn_metadata: AttentionMetadata,
    residual: Optional[torch.Tensor] = None,
    smooth_quant_scale: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    qkv, _ = self.qkv_proj(hidden_states, smooth_quant_scale)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: pack q & k to fit tmo.apply_rotary
    @optimization: avoid redundant split operation
    '''
    if self.rope_scaling is not None and self.rope_scaling["rope_type"] == "longrope":
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q, k = self.rotary_emb(positions, q, k)
    else:
        # Optimized: split qkv into [qk, v] directly, avoiding redundant split
        qk, v = 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))
        # Split qk into q and k after rotary embedding
        q, k = qk.split([self.q_size, self.kv_size], dim=-1)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    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


def vllm__module_executor__models__llama__LlamaDecoderLayer____init__(
    self,
    config: LlamaConfig,
    cache_config: Optional[CacheConfig] = None,
    quant_config: Optional[QuantizationConfig] = None,
    prefix: str = "",
) -> None:
    super(LlamaDecoderLayer, self).__init__()
    self.hidden_size = config.hidden_size
    rope_theta = getattr(config, "rope_theta", 10000)
    rope_scaling = getattr(config, "rope_scaling", None)
    if rope_scaling is not None and getattr(
            config, "original_max_position_embeddings", None):
        rope_scaling["original_max_position_embeddings"] = (
            config.original_max_position_embeddings)
    max_position_embeddings = getattr(config, "max_position_embeddings",
                                        8192)
    # Support abacusai/Smaug-72B-v0.1 with attention_bias
    # Support internlm/internlm-7b with bias
    attention_bias = getattr(config, "attention_bias", False) or getattr(
        config, "bias", False)
    self.self_attn = LlamaAttention(
        config=config,
        hidden_size=self.hidden_size,
        num_heads=config.num_attention_heads,
        num_kv_heads=getattr(config, "num_key_value_heads",
                                config.num_attention_heads),
        rope_theta=rope_theta,
        rope_scaling=rope_scaling,
        max_position_embeddings=max_position_embeddings,
        quant_config=quant_config,
        bias=attention_bias,
        cache_config=cache_config,
        prefix=f"{prefix}.self_attn",
    )
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: use FeedForward instead of MLP
    '''
    self.mlp = FeedForward(hidden_size=config.hidden_size,
                            intermediate_size=config.intermediate_size,
                            hidden_act='silu',
                            up_proj_name='gate_up_proj',
                            is_gated=True,
                            down_proj_name='down_proj',
                            bias=getattr(config, "mlp_bias", False),
                            quant_config=quant_config,
                            prefix=f"{prefix}.mlp")
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    self.input_layernorm = RMSNorm(config.hidden_size,
                                    eps=config.rms_norm_eps)
    self.post_attention_layernorm = RMSNorm(config.hidden_size,
                                            eps=config.rms_norm_eps)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: prepare to perf sq cases if suitable
    '''
    self.is_per_tesnor_sq_perf_cases = is_per_tensor_smoothquant(quant_config)
    self.is_per_token_sq_perf_cases = is_per_token_smoothquant(quant_config)
    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
        self.self_attn.qkv_proj.quant_method.skip_quant_input = True
        self.mlp.gate_up_proj.quant_method.skip_quant_input = True
        self.quant_fusion_attn_layernorm = None
        self.quant_fusion_mlp_layernorm = None
    '''
    ==================
    End of MLU Hijack
    ==================
    '''

def vllm__module_executor__models__llama__LlamaDecoderLayer__forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
    ) -> torch.Tensor:
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: perf model by:
    1) add residual in matmul;
    2) fuse quantization in layernorm in per-tensor sq case;
    '''
    attn_layernorm = self.input_layernorm
    mlp_layernorm = self.post_attention_layernorm
    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
        if self.quant_fusion_attn_layernorm is None:
            if self.is_per_token_sq_perf_cases:
                attn_quant_scale = self.self_attn.qkv_proj.smooth
                mlp_quant_scale = self.mlp.gate_up_proj.smooth
            else:
                attn_quant_scale = self.self_attn.qkv_proj.scale_to_int
                mlp_quant_scale = self.mlp.gate_up_proj.scale_to_int
            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
                self.input_layernorm, attn_quant_scale,
                dynamic_quant=self.is_per_token_sq_perf_cases)
            self.quant_fusion_mlp_layernorm = quant_fusion_with_rmsnorm(
                self.post_attention_layernorm, mlp_quant_scale,
                dynamic_quant=self.is_per_token_sq_perf_cases)
        attn_layernorm = self.quant_fusion_attn_layernorm
        mlp_layernorm = self.quant_fusion_mlp_layernorm

    return decoder_layer_forward_base(positions, hidden_states, kv_cache, attn_metadata,
                                      attn_layernorm,
                                      self.self_attn,
                                      mlp_layernorm,
                                      self.mlp,
                                      input_norm_fuse_en=self.is_per_token_sq_perf_cases,
                                      post_norm_fuse_en=self.is_per_token_sq_perf_cases)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''

def vllm__module_executor__models__llama__LlamaModel__forward(
        self,
        input_ids: Optional[torch.Tensor],
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors],
        inputs_embeds: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
    '''
    =============================
    Modify by vllm_mlu
    =============================
    '''
    return decoder_model_forward_base_pp(input_ids, positions, kv_caches, attn_metadata, intermediate_tensors,
                                         self.layers, self.start_layer, self.end_layer,
                                         self.get_input_embeddings,
                                         self.norm,
                                         inputs_embeds)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''

MluHijackObject.apply_hijack(LlamaAttention,
                             LlamaAttention.__init__,
                             vllm__module_executor__models__llama__LlamaAttention____init__)
MluHijackObject.apply_hijack(LlamaAttention,
                             LlamaAttention.forward,
                             vllm__module_executor__models__llama__LlamaAttention__forward)
MluHijackObject.apply_hijack(LlamaDecoderLayer,
                             LlamaDecoderLayer.__init__,
                             vllm__module_executor__models__llama__LlamaDecoderLayer____init__)
MluHijackObject.apply_hijack(LlamaDecoderLayer,
                             LlamaDecoderLayer.forward,
                             vllm__module_executor__models__llama__LlamaDecoderLayer__forward)
MluHijackObject.apply_hijack(LlamaModel,
                             LlamaModel.forward,
                             vllm__module_executor__models__llama__LlamaModel__forward)
