from typing import Optional

import torch
import torch.nn as nn
from transformers import CLIPVisionConfig

from vllm.distributed import get_tensor_model_parallel_world_size
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm_mlu.model_executor.layers.feed_forward import FeedForward
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.models.clip import (CLIPVisionModel,
                                             CLIPVisionTransformer,
                                             CLIPEncoderLayer,
                                             CLIPParallelAttention)
from vllm.model_executor.layers.linear import (QKVParallelLinear,
                                               RowParallelLinear)
from vllm.logger import init_logger

logger = init_logger(__name__)


class MLUCLIPAttention(nn.Module):
    """
    MLU-compatible CLIP attention implementation.
    Used as fallback when num_heads % tp_size != 0.
    This implementation does not use tensor parallelism.
    """

    def __init__(
        self,
        config: CLIPVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.embed_dim // self.num_heads

        if self.head_dim * self.num_heads != self.embed_dim:
            raise ValueError(
                f"embed_dim must be divisible by num_heads "
                f"(got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})."
            )

        self.scale = self.head_dim ** -0.5
        self.dropout = getattr(config, 'attention_dropout', 0.0)

        # Use non-parallel linear layers since this is fallback for non-divisible cases
        self.qkv_proj = QKVParallelLinear(
            hidden_size=self.embed_dim,
            head_size=self.head_dim,
            total_num_heads=self.num_heads,
            quant_config=quant_config,
            prefix=f"{prefix}.qkv_proj",
            bias=True,
        )

        self.out_proj = RowParallelLinear(
            input_size=self.embed_dim,
            output_size=self.embed_dim,
            quant_config=quant_config,
            prefix=f"{prefix}.out_proj",
            bias=True,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
    ):
        """
        Input shape: Batch x Time x Channel
        Compatible with CLIPSdpaAttention interface
        """
        bsz, tgt_len, _ = hidden_states.size()

        # Project to Q, K, V
        qkv_states, _ = self.qkv_proj(hidden_states)
        query_states, key_states, value_states = qkv_states.chunk(3, dim=-1)

        # Reshape for attention computation
        query_states = query_states.view(bsz, tgt_len, self.num_heads, self.head_dim)
        key_states = key_states.view(bsz, tgt_len, self.num_heads, self.head_dim)
        value_states = value_states.view(bsz, tgt_len, self.num_heads, self.head_dim)

        # Use MLU flash attention for inference
        if self.dropout == 0.0:
            try:
                from vllm import _mlu_ops as mlu_ops
                out = mlu_ops.flash_attention(
                    query_states,
                    key_states,
                    value_states,
                    out=None,
                    cu_seq_lens_q=None,
                    cu_seq_lens_kv=None,
                    alibi_slope=None,
                    attn_bias=None,
                    max_seq_len_q=tgt_len,
                    max_seq_len_kv=tgt_len,
                    softmax_scale=self.scale,
                    is_causal=False
                )
            except (ImportError, AttributeError):
                # Fallback to standard PyTorch attention if MLU ops not available
                logger.warning("MLU ops not available, using standard PyTorch attention")
                out = self._pytorch_attention(query_states, key_states, value_states)
        else:
            # Use xformers if dropout is needed (training mode)
            try:
                from xformers import ops as xops
                out = xops.memory_efficient_attention_forward(
                    query_states,
                    key_states,
                    value_states,
                    p=self.dropout,
                    scale=self.scale
                )
            except ImportError:
                logger.warning("xformers not available, using standard PyTorch attention")
                out = self._pytorch_attention(query_states, key_states, value_states)

        # Reshape output
        out = out.view(bsz, tgt_len, -1)

        # Output projection
        attn_output, _ = self.out_proj(out)

        return attn_output, None

    def _pytorch_attention(
        self,
        query_states: torch.Tensor,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
    ) -> torch.Tensor:
        """
        Standard PyTorch scaled dot-product attention as fallback.
        Input shape: [batch, seq_len, num_heads, head_dim]
        """
        # Transpose to [batch, num_heads, seq_len, head_dim]
        query_states = query_states.transpose(1, 2)
        key_states = key_states.transpose(1, 2)
        value_states = value_states.transpose(1, 2)

        # Compute attention scores
        attn_weights = torch.matmul(query_states, key_states.transpose(-2, -1)) * self.scale
        attn_weights = torch.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)

        # Apply attention to values
        attn_output = torch.matmul(attn_weights, value_states)

        # Transpose back to [batch, seq_len, num_heads, head_dim]
        attn_output = attn_output.transpose(1, 2)

        return attn_output


def vllm__module_executor__models__clip__CLIPParallelAttention__forward(
    self,
    hidden_states: torch.Tensor,
    residual: torch.Tensor
):
    """Input shape: Batch x Time x Channel"""
    bsz, tgt_len, _ = hidden_states.size()

    qkv_states, _ = self.qkv_proj(hidden_states)
    query_states, key_states, value_states = qkv_states.chunk(3, dim=-1)

    query_states = query_states.view(bsz, tgt_len,
                                     self.num_heads_per_partition,
                                     self.head_dim)
    key_states = key_states.view(bsz, tgt_len,
                                 self.num_heads_per_partition,
                                 self.head_dim)
    value_states = value_states.view(bsz, tgt_len,
                                     self.num_heads_per_partition,
                                     self.head_dim)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: perf attn using tmo flash attn
    '''
    if self.dropout is None or self.dropout == 0.0:
        # Always true for inference
        from vllm import _mlu_ops as mlu_ops

        out = mlu_ops.flash_attention(query_states,
                                      key_states,
                                      value_states,
                                      out=None,
                                      cu_seq_lens_q=None,
                                      cu_seq_lens_kv=None,
                                      alibi_slope=None,
                                      attn_bias=None,
                                      max_seq_len_q=tgt_len,
                                      max_seq_len_kv=tgt_len,
                                      softmax_scale=self.scale,
                                      is_causal=False)
    else:
        from xformers import ops as xops

        out = xops.memory_efficient_attention_forward(query_states,
                                                      key_states,
                                                      value_states,
                                                      p=self.dropout,
                                                      scale=self.scale)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    out = out.view(bsz, tgt_len, -1)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: add residual
    ''' 
    attn_output, _ = self.out_proj(out, residual)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    return attn_output, None


def vllm__module_executor__models__clip__CLIPEncoderLayer____init__(
    self,
    config: CLIPVisionConfig,
    quant_config: Optional[QuantizationConfig] = None,
    prefix: str = "",
):
    super(CLIPEncoderLayer, self).__init__()

    num_heads = config.num_attention_heads
    tp_size = get_tensor_model_parallel_world_size()
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: perf attn using tmo flash attn, do not check xformers
    '''
    if num_heads % tp_size == 0:
        self.self_attn = CLIPParallelAttention(
            config,
            quant_config=quant_config,
            prefix=f"{prefix}.self_attn",
        )
        self.use_parallel_attn = True
    else:
        logger.warning("Use MLUCLIPAttention for clip model (fallback for non-divisible heads).")
        self.self_attn = MLUCLIPAttention(
            config,
            quant_config=quant_config,
            prefix=f"{prefix}.self_attn",
        )
        self.use_parallel_attn = False
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    self.layer_norm1 = nn.LayerNorm(config.hidden_size,
                                    eps=config.layer_norm_eps)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: use FeedForward instead of MLP
    '''
    self.mlp = FeedForward(hidden_size=config.hidden_size,
                           intermediate_size=config.intermediate_size,
                           is_gated=False,
                           bias=True,
                           hidden_act=config.hidden_act,
                           quant_config=quant_config,
                           up_proj_name='fc1',
                           down_proj_name='fc2',
                           prefix=f"{prefix}.mlp")
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    self.layer_norm2 = nn.LayerNorm(config.hidden_size,
                                    eps=config.layer_norm_eps)


def vllm__module_executor__models__clip__CLIPEncoderLayer__forward(
    self,
    hidden_states: torch.Tensor
) -> torch.Tensor:
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: apply residual fusion
    '''
    residual = hidden_states
    if self.use_parallel_attn:
        hidden_states = self.layer_norm1(hidden_states)
        hidden_states, _ = self.self_attn(hidden_states, residual)
    else:
        hidden_states = self.layer_norm1(hidden_states)
        hidden_states, _ = self.self_attn(hidden_states)
        hidden_states = residual + hidden_states

    residual = hidden_states
    hidden_states = self.layer_norm2(hidden_states)
    bsz, tgt_len, _ = hidden_states.size()
    hidden_states = self.mlp(hidden_states, residual)
    hidden_states = hidden_states.view(bsz, tgt_len, -1)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    return hidden_states


def vllm__module_executor__models__clip__CLIPVisionModel____init__(
    self,
    config: CLIPVisionConfig,
    quant_config: Optional[QuantizationConfig] = None,
    *,
    num_hidden_layers_override: Optional[int] = None,
    require_post_norm: Optional[bool] = None,
    prefix: str = "",
) -> None:
    super(CLIPVisionModel, self).__init__()

    tp_size = get_tensor_model_parallel_world_size()
    num_heads = config.num_attention_heads
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: perf attn using tmo flash attn, do not check xformers
    '''
    self.shard_weight = num_heads % tp_size == 0
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    self.vision_model = CLIPVisionTransformer(
        config=config,
        quant_config=quant_config,
        num_hidden_layers_override=num_hidden_layers_override,
        require_post_norm=require_post_norm,
        prefix=f"{prefix}.vision_model",
    )


MluHijackObject.apply_hijack(CLIPParallelAttention,
                             CLIPParallelAttention.forward,
                             vllm__module_executor__models__clip__CLIPParallelAttention__forward)
MluHijackObject.apply_hijack(CLIPEncoderLayer,
                             CLIPEncoderLayer.__init__,
                             vllm__module_executor__models__clip__CLIPEncoderLayer____init__)
MluHijackObject.apply_hijack(CLIPEncoderLayer,
                             CLIPEncoderLayer.forward,
                             vllm__module_executor__models__clip__CLIPEncoderLayer__forward)
MluHijackObject.apply_hijack(CLIPVisionModel,
                             CLIPVisionModel.__init__,
                             vllm__module_executor__models__clip__CLIPVisionModel____init__)
