#
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
# This file is a part of the vllm-ascend project.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from collections.abc import Callable

import torch
from vllm.distributed import get_dp_group, get_ep_group, get_tp_group
from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig

from vllm_ascend.ascend_forward_context import _EXTRA_CTX, MoECommType
from vllm_ascend.ops.fused_moe.experts_selector import zero_experts_compute
from vllm_ascend.ops.fused_moe.fused_moe import AscendMoERunner
from vllm_ascend.ops.fused_moe.moe_comm_method import (
    AllGatherCommImpl,
    FusedExpertsResult,
    _MoECommMethods,
)
from vllm_ascend.ops.fused_moe.moe_runtime_args import build_fused_experts_input
from vllm_ascend.quantization.quant_type import QuantType
from vllm_ascend.utils import maybe_trans_nz, vllm_version_is

from .experts_selector import select_experts
from .moe_comm_method import AllGatherCommImpl310

if vllm_version_is("0.23.0"):
    from vllm.model_executor.layers.fused_moe.layer import FusedMoE as _LegacyFusedMoEBase
    from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod
else:
    from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import UnquantizedFusedMoEMethod

    try:
        from vllm.model_executor.layers.fused_moe.layer import FusedMoE as _LegacyFusedMoEBase
    except ImportError:
        _LegacyFusedMoEBase = torch.nn.Module

    if not isinstance(_LegacyFusedMoEBase, type):
        _LegacyFusedMoEBase = torch.nn.Module


class AscendUnquantizedFusedMoEMethod310(UnquantizedFusedMoEMethod):
    def __init__(self, moe: FusedMoEConfig = None):
        super().__init__(moe=moe)

    @property
    def is_monolithic(self) -> bool:
        return False

    def maybe_make_prepare_finalize(self, routing_tables=None):
        # Ascend 310P uses its own MoE communication and forward_impl path.
        # Do not let upstream modular-kernel initialization replace it.
        return None

    def process_weights_after_loading(self, layer):
        super().process_weights_after_loading(layer)

        # Fused gate_up_proj (column parallel)
        w13_data = self._maybe_pad_weight(layer.w13_weight.data).transpose(1, 2).contiguous()
        w13_data = maybe_trans_nz(w13_data)
        layer.w13_weight = torch.nn.Parameter(w13_data, requires_grad=False)
        # down_proj (row parallel)
        w2_data = self._maybe_pad_weight(layer.w2_weight.data).transpose(1, 2).contiguous()
        w2_data = maybe_trans_nz(w2_data)
        layer.w2_weight = torch.nn.Parameter(w2_data, requires_grad=False)

    def apply(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        use_grouped_topk: bool,
        top_k: int,
        router_logits: torch.Tensor,
        renormalize: bool,
        topk_group: int | None = None,
        num_expert_group: int | None = None,
        custom_routing_function: Callable | None = None,
        scoring_func: str = "softmax",
        e_score_correction_bias: torch.Tensor | None = None,
        num_experts: int = -1,
        expert_map: torch.Tensor | None = None,
        apply_router_weight_on_input: bool = False,
        **kwargs,
    ) -> torch.Tensor:
        zero_expert_num = getattr(layer, "zero_expert_num", 0)
        zero_expert_type = getattr(layer, "zero_expert_type", None)

        topk_weights, topk_ids = select_experts(
            hidden_states=x,
            router_logits=router_logits,
            top_k=top_k,
            use_grouped_topk=use_grouped_topk,
            renormalize=renormalize,
            topk_group=topk_group,
            num_expert_group=num_expert_group,
            custom_routing_function=custom_routing_function,
            scoring_func=scoring_func,
            e_score_correction_bias=e_score_correction_bias,
            global_num_experts=num_experts,
        )

        if zero_expert_num > 0 and zero_expert_type is not None:
            topk_ids, topk_weights, zero_expert_result = zero_experts_compute(
                expert_indices=topk_ids,
                expert_scales=topk_weights,
                num_experts=num_experts,
                zero_expert_type=zero_expert_type,
                hidden_states=x,
            )

        topk_weights = topk_weights.to(x.dtype)

        moe_comm_method = _EXTRA_CTX.moe_comm_method
        final_hidden_states = moe_comm_method.fused_experts(
            fused_experts_input=build_fused_experts_input(
                hidden_states=x,
                topk_weights=topk_weights,
                topk_ids=topk_ids,
                w1=layer.w13_weight,
                w2=layer.w2_weight,
                quant_type=QuantType.NONE,
                dynamic_eplb=False,
                expert_map=expert_map,
                apply_router_weight_on_input=apply_router_weight_on_input,
            ),
        )
        if zero_expert_num > 0 and zero_expert_type is not None:
            final_hidden_states += zero_expert_result
        return final_hidden_states


if not vllm_version_is("0.23.0"):

    class AscendMoERunner310(AscendMoERunner):
        def __init__(
            self,
            layer_name,
            moe_config,
            router,
            routed_experts,
            enable_dbo=False,
            gate=None,
            shared_experts=None,
            shared_expert_gate=None,
            routed_input_transform=None,
            routed_output_transform=None,
            routed_scaling_factor=1,
            tid2eid=None,
            n_shared_experts: int = 0,
        ):
            super().__init__(
                layer_name,
                moe_config,
                router,
                routed_experts,
                enable_dbo,
                gate,
                shared_experts,
                shared_expert_gate,
                routed_input_transform,
                routed_output_transform,
                routed_scaling_factor,
                tid2eid,
                n_shared_experts,
            )

            if routed_experts.quant_config is None:
                routed_experts.quant_method = AscendUnquantizedFusedMoEMethod310(self.moe_config)
                self.quant_type = self._get_quant_type()

            self.multistream_overlap_gate = False
            self.shared_multistream_overlap_gate = False
            self.multistream_overlap_shared_expert = False
            _MoECommMethods[MoECommType.ALLGATHER] = AllGatherCommImpl310(self.moe_config)


class AscendFusedMoE310(_LegacyFusedMoEBase):
    def __init__(self, *args, **kwargs):
        if _LegacyFusedMoEBase is torch.nn.Module:
            raise RuntimeError("AscendFusedMoE310 is only kept for the legacy FusedMoE class API.")
        super().__init__(*args, **kwargs)

        self._routed_input_transform = kwargs.get("routed_input_transform")
        self._shared_experts = kwargs.get("shared_experts")
        self.global_num_experts = kwargs["num_experts"]

        if self.quant_config is None:
            self.quant_method = AscendUnquantizedFusedMoEMethod310(self.moe_config)
        else:
            self.quant_method = self.quant_config.get_quant_method(self, self.layer_name)

        assert self.quant_method is not None
        # Keep base_quant_method aligned with the Ascend-replaced quant_method
        # so FusedMoE.maybe_init_modular_kernel doesn't dispatch into the
        # upstream UnquantizedFusedMoEMethod.maybe_make_prepare_finalize.
        self.base_quant_method = self.quant_method

        self.moe_config.tp_group = get_tp_group()
        self.moe_config.dp_group = get_dp_group()
        self.moe_config.ep_group = get_ep_group()
        self.moe_config.supports_eplb = False

        # init moe
        self.global_expert_map = None
        self.local_expert_map = None
        if self.moe_config.ep_size > 1:
            raise RuntimeError("Expert Parallel is not supported on 310P. Please remove --enable-expert-parallel.")
        self.local_num_experts = self.global_num_experts

        self.moe_config.num_experts = self.global_num_experts
        self.moe_config.num_local_experts = self.local_num_experts
        self.moe_config.global_redundant_expert_num = 0

        moe_quant_params = {
            "num_experts": self.local_num_experts,
            "hidden_size": self.hidden_size,
            "intermediate_size_per_partition": self.intermediate_size_per_partition,
            "params_dtype": self.params_dtype,
            "weight_loader": self.weight_loader,
        }

        self.quant_method.create_weights(layer=self, **moe_quant_params)
        self.quant_type = self.get_quant_type()

        _MoECommMethods[MoECommType.ALLGATHER] = AllGatherCommImpl310(self.moe_config)

        if vllm_version_is("0.23.0"):
            self.runner = AscendMoERunner(
                self.layer_name,
                self.moe_config,
                self.router,
                self._routed_input_transform,
                kwargs.pop("gate", None),
                kwargs.pop("shared_experts", None),
                self.quant_method,
                self.vllm_config.parallel_config.enable_dbo,
            )
        else:
            self.runner = AscendMoERunner310(
                self.layer_name,
                self.moe_config,
                self.router,
                self._routed_input_transform,
                kwargs.pop("gate", None),
                kwargs.pop("shared_experts", None),
                self.quant_method,
                self.vllm_config.parallel_config.enable_dbo,
            )

    @property
    def is_internal_router(self) -> bool:
        # 310P Ascend path expects router logits from the model forward path.
        return False

    def init_experts_map(self, moe_config):
        """
        Initialize expert mapping for MoE (Mixture of Experts) model.

        This function creates mappings between global expert indices and local expert indices
        for each rank in the expert parallel group. It divides the total experts among
        different ranks and creates both global and local expert maps that are used
        during MoE computation to determine which experts are handled by which rank.

        Args:
            moe_config: Configuration object containing MoE parameters including
                       number of experts, expert parallel size, and expert parallel rank.

        Returns:
            tuple: A tuple containing:
                   - global_expert_map: Stack of expert maps for all ranks
                   - local_expert_map: Expert map for the current rank (transferred to NPU)
        """
        n_experts = moe_config.num_experts
        ep_size = moe_config.ep_size
        all_experts = torch.arange(n_experts, dtype=torch.int32)
        experts_groups = all_experts.chunk(ep_size)
        global_expert_map = []
        local_expert_map = None
        for rankid in range(ep_size):
            expert_map = torch.full((n_experts,), -1, dtype=torch.int32)
            local_experts = experts_groups[rankid]
            expert_map[local_experts] = torch.arange(local_experts.shape[0], dtype=torch.int32)
            global_expert_map.append(expert_map)
            if rankid == moe_config.ep_rank:
                local_expert_map = expert_map.npu()
        return torch.stack(global_expert_map), local_expert_map

    def get_quant_type(self) -> QuantType:
        quant_method = self.quant_method
        if not hasattr(quant_method, "quant_method") or quant_method.quant_method is None:
            return QuantType.NONE

        method = quant_method.quant_method
        quant_type = getattr(method, "quant_type", QuantType.NONE)
        if quant_type not in [QuantType.NONE, QuantType.W8A8]:
            raise RuntimeError("Only Unquant and W8A8 is supported.")
        return quant_type

    def forward_impl(  # type: ignore[override]
        self, hidden_states: torch.Tensor, router_logits: torch.Tensor
    ) -> torch.Tensor:
        assert self.quant_method is not None
        assert self.routed_scaling_factor == 1.0, "routed_scaling_factor != 1.0 is not supported."

        prepare_output = _EXTRA_CTX.moe_comm_method.prepare(
            hidden_states=hidden_states, router_logits=router_logits, quant_type=self.quant_type
        )
        hidden_states = prepare_output.hidden_states
        router_logits = prepare_output.router_logits
        pertoken_scale = prepare_output.pertoken_scale
        padded_hidden_states_shape = prepare_output.padded_hidden_states_shape

        # Matrix multiply.
        fused_experts_results: FusedExpertsResult = self.quant_method.apply(
            layer=self,
            x=hidden_states,
            use_grouped_topk=self.use_grouped_topk,
            top_k=self.top_k,
            router_logits=router_logits,
            renormalize=self.renormalize,
            topk_group=self.topk_group,
            num_expert_group=self.num_expert_group,
            custom_routing_function=self.custom_routing_function,
            scoring_func=self.scoring_func,
            e_score_correction_bias=self.e_score_correction_bias,
            num_experts=self.global_num_experts,
            expert_map=self.local_expert_map,
            apply_router_weight_on_input=self.apply_router_weight_on_input,
            pertoken_scale=pertoken_scale,
        )

        routed_out = _EXTRA_CTX.moe_comm_method.finalize(
            hidden_states=fused_experts_results.routed_out,
            reduce_results=isinstance(_EXTRA_CTX.moe_comm_method, AllGatherCommImpl),
            padded_hidden_states_shape=padded_hidden_states_shape,
        )

        return routed_out

    def _forward_shared_experts(self, hidden_states: torch.Tensor):
        if self._shared_experts is None:
            return None
        return self._shared_experts(hidden_states)

    def shared_forward_impl(  # type: ignore[override]
        self, hidden_states: torch.Tensor, router_logits: torch.Tensor
    ):
        routed_out = AscendFusedMoE310.forward_impl(
            self,
            hidden_states=hidden_states,
            router_logits=router_logits,
        )
        if self._shared_experts is None:
            return routed_out
        shared_out = self._forward_shared_experts(hidden_states)
        return shared_out, routed_out