import torch
from typing import Callable, Optional, List, Union, Tuple

from vllm import _mlu_ops as mlu_ops
from vllm.attention import AttentionMetadata
from vllm.sequence import IntermediateTensors
from vllm.distributed import get_pp_group
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from transformers import PretrainedConfig
from vllm_mlu._mlu_utils import check_context_comm_cmpt_parallel

def hunyuan_decoder_layer_forward_base(
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
        input_layernorm: Callable,
        self_attn: Callable,
        post_layernorm: Callable,
        mlp: Callable,
        kv_states: Optional[Tuple[torch.Tensor]] = None,
        apply_residual_connection_post_layernorm: bool = False,
        position_name: str = 'positions',
        input_norm_fuse_en: bool = False,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
    smooth_quant_scale = None
    if input_norm_fuse_en:
        layernorm_output, smooth_quant_scale = input_layernorm(hidden_states)
    else:
        layernorm_output = input_layernorm(hidden_states)
    if apply_residual_connection_post_layernorm:
        residual = layernorm_output
    else:
        residual = hidden_states
    
    # Self Attention
    attention_output, ori_kv_states = self_attn(
        **{position_name: positions},
        hidden_states=layernorm_output,
        kv_cache=kv_cache,
        attn_metadata=attn_metadata,
        residual=residual,
        kv_states=kv_states,
        smooth_quant_scale=smooth_quant_scale,
    )
    
    layernorm_output = post_layernorm(attention_output)
    if apply_residual_connection_post_layernorm:
        residual = layernorm_output
    else:
        residual = attention_output

    # Fully Connected      
    hidden_states = mlp(layernorm_output, residual)
    return hidden_states, ori_kv_states

def decoder_layer_forward_base(
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
        input_layernorm: Callable,
        self_attn: Callable,
        post_layernorm: Callable,
        mlp: Callable,
        apply_residual_connection_post_layernorm: bool = False,
        position_name: str = 'positions',
        input_norm_fuse_en: bool = False,
        post_norm_fuse_en: bool = False,
    ) -> torch.Tensor:
    smooth_quant_scale = None
    if input_norm_fuse_en:
        layernorm_output, smooth_quant_scale = input_layernorm(hidden_states)
    else:
        layernorm_output = input_layernorm(hidden_states)

    if apply_residual_connection_post_layernorm:
        residual = layernorm_output
    else:
        residual = hidden_states
    
    # Self Attention
    attention_output = self_attn(
        **{position_name: positions},
        hidden_states=layernorm_output,
        kv_cache=kv_cache,
        attn_metadata=attn_metadata,
        residual=residual,
        smooth_quant_scale=smooth_quant_scale,
    )
    smooth_quant_scale = None
    if post_norm_fuse_en:
        layernorm_output, smooth_quant_scale = post_layernorm(attention_output)
    else:
        layernorm_output = post_layernorm(attention_output)
    if apply_residual_connection_post_layernorm:
        residual = layernorm_output
    else:
        residual = attention_output

    # Fully Connected      
    kwargs = dict()
    if post_norm_fuse_en:
        kwargs['smooth_quant_scale'] = smooth_quant_scale
    hidden_states = mlp(layernorm_output, residual, **kwargs)
    return hidden_states


def decoder_model_forward_base(
        input_ids: Optional[torch.Tensor],
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        layers: torch.nn.ModuleList,
        get_input_embeddings: Callable,
        norm: Callable
    ) -> torch.Tensor:
    hidden_states = get_input_embeddings(input_ids)
    for i in range(len(layers)):
        layer = layers[i]
        hidden_states = layer(
            positions,
            hidden_states,
            kv_caches[i],
            attn_metadata,
        )
    hidden_states = norm(hidden_states)
    return hidden_states


def hunyuan_decoder_model_forward_base_pp(
        config: PretrainedConfig,
        input_ids: Optional[torch.Tensor],
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors],
        layers: torch.nn.ModuleList,
        start_layer: int,
        end_layer: int,
        get_input_embeddings: Callable,
        norm: Callable,
        inputs_embeds: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
    if get_pp_group().is_first_rank:
        if inputs_embeds is not None:
            hidden_states = inputs_embeds
        else:
            hidden_states = get_input_embeddings(input_ids)
    else:
        assert intermediate_tensors is not None
        hidden_states = intermediate_tensors["hidden_states"]

    cla_factor = getattr(config, "cla_share_factor", 1)
    prev_kv_states = None
    for i in range(start_layer, end_layer):
        layer = layers[i]
        hidden_states, kv_states = layer(
            positions,
            hidden_states,
            kv_caches[i - start_layer],
            attn_metadata,
            prev_kv_states,
        )
        if (i - start_layer) % cla_factor == 0:
            prev_kv_states = kv_states
        else:
            prev_kv_states = None

    if not get_pp_group().is_last_rank:
        return IntermediateTensors({
            "hidden_states": hidden_states,
        })

    hidden_states = norm(hidden_states)
    return hidden_states

def decoder_model_forward_base_pp(
        input_ids: Optional[torch.Tensor],
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors],
        layers: torch.nn.ModuleList,
        start_layer: int,
        end_layer: int,
        get_input_embeddings: Callable,
        norm: Callable,
        inputs_embeds: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
    if get_pp_group().is_first_rank:
        if inputs_embeds is not None:
            hidden_states = inputs_embeds
        else:
            hidden_states = get_input_embeddings(input_ids)
    else:
        assert intermediate_tensors is not None
        hidden_states = intermediate_tensors["hidden_states"]

    for i in range(start_layer, end_layer):
        layer = layers[i]
        hidden_states = layer(
            positions,
            hidden_states,
            kv_caches[i - start_layer],
            attn_metadata,
        )

    if not get_pp_group().is_last_rank:
        return IntermediateTensors({
            "hidden_states": hidden_states,
        })

    hidden_states = norm(hidden_states)
    return hidden_states


def is_smoothquant(quant_config: QuantizationConfig) -> bool:
    return (quant_config is not None and quant_config.get_name() == "SmoothQuant")

def is_per_tensor_smoothquant(quant_config: QuantizationConfig) -> bool:
    return is_smoothquant(quant_config) and quant_config.input_quant_method == "per_tensor"

def is_per_token_smoothquant(quant_config: QuantizationConfig) -> bool:
    if check_context_comm_cmpt_parallel():
        return False

    return is_smoothquant(quant_config) and quant_config.input_quant_method == "per_token"

def quant_fusion_with_layernorm(
        op: torch.nn.LayerNorm,
        quant_scale: torch.Tensor,
        dynamic_quant: bool = False,
    ) -> Callable:
    bias = None
    if op.bias is not None:
        bias = op.bias.data

    def func(x: torch.Tensor) -> torch.Tensor:
        return mlu_ops.fused_layer_norm(
            x,
            None,
            op.weight.data,
            bias,
            None,
            op.eps,
            False,
            quant_scale,
            dynamic_quant)

    return func


def quant_fusion_with_rmsnorm(
        op: RMSNorm,
        quant_scale: torch.Tensor,
        dynamic_quant: bool = False,
    ) -> Callable:

    def func(x: torch.Tensor) -> torch.Tensor:
        return mlu_ops.fused_rms_norm(
            x,
            None,
            op.weight.data,
            None,
            None,
            op.variance_epsilon,
            False,
            quant_scale,
            dynamic_quant)

    return func
