# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Optional

import torch

import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
from vllm.forward_context import get_forward_context


class DeepEPAutoPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
    """
    Auto Prepare/Finalize that wraps both DeepEP High-Throughput and
    Low-Latency implementations and selects one based on prefill/decode phase.
    """

    def __init__(self,
                 ht_prepare_finalize: mk.FusedMoEPrepareAndFinalize,
                 ll_prepare_finalize: mk.FusedMoEPrepareAndFinalize):
        super().__init__()
        self.ht_prepare_finalize = ht_prepare_finalize
        self.ll_prepare_finalize = ll_prepare_finalize
        self._current_phase = "decode"  # default to decode (LL)

    def _get_current_prepare_finalize(self) -> mk.FusedMoEPrepareAndFinalize:
        """Get the appropriate prepare_finalize based on current phase."""
        # Try to infer phase from forward_context if available:
        # try:
        #     forward_context = get_forward_context()
        #     attn_metadata = forward_context.attn_metadata
        #     # Handle both v0 (single AttentionMetadata) and v1 (dict) formats
        #     if isinstance(attn_metadata, dict):
        #         if attn_metadata:
        #             attn_metadata = next(iter(attn_metadata.values()))
        #         else:
        #             attn_metadata = None
            
        #     if attn_metadata is not None and hasattr(attn_metadata,
        #                                              "num_decode_tokens"):
        #         # 只根据 decode tokens 判定：有 decode -> decode，否则 prefill
        #         self._current_phase = ("decode"
        #                                if attn_metadata.num_decode_tokens > 0
        #                                else "prefill")
        # except Exception:
        #     # If forward_context is not available, use stored phase
        #     pass

        # Prefill uses HT, decode uses LL
        if self._current_phase == "prefill":
            #rint("************prefill***********")
            return self.ll_prepare_finalize
        else:
            # print("attn_metadata.num_decode_tokens",attn_metadata.num_decode_tokens)
            return self.ht_prepare_finalize
        #return self.ht_prepare_finalize
    @property
    def activation_format(self) -> mk.FusedMoEActivationFormat:
        pf = self._get_current_prepare_finalize()
        try:
            return pf.activation_format
        except NotImplementedError:
            # Fallback to standard format if underlying impl does not provide it.
            return mk.FusedMoEActivationFormat.Standard

    def topk_indices_dtype(self) -> Optional[torch.dtype]:
        pf = self._get_current_prepare_finalize()
        return pf.topk_indices_dtype()

    def max_num_tokens_per_rank(self) -> Optional[int]:
        pf = self._get_current_prepare_finalize()
        return pf.max_num_tokens_per_rank()

    def num_dispatchers(self) -> int:
        pf = self._get_current_prepare_finalize()
        return pf.num_dispatchers()

    def prepare_async(
        self,
        a1: torch.Tensor,
        a1_scale: Optional[torch.Tensor],
        a2_scale: Optional[torch.Tensor],
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        num_experts: int,
        expert_map: Optional[torch.Tensor],
        apply_router_weight_on_input: bool,
        quant_config: FusedMoEQuantConfig,
    ):
        pf = self._get_current_prepare_finalize()
        return pf.prepare_async(
            a1, a1_scale, a2_scale, topk_weights, topk_ids,
            num_experts, expert_map, apply_router_weight_on_input, quant_config)

    def prepare(
        self,
        a1: torch.Tensor,
        a1_scale: Optional[torch.Tensor],
        a2_scale: Optional[torch.Tensor],
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        num_experts: int,
        expert_map: Optional[torch.Tensor],
        apply_router_weight_on_input: bool,
        quant_config: FusedMoEQuantConfig,
    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],
               Optional[torch.Tensor], Optional[torch.Tensor]]:
        pf = self._get_current_prepare_finalize()
        return pf.prepare(
            a1, a1_scale, a2_scale, topk_weights, topk_ids,
            num_experts, expert_map, apply_router_weight_on_input, quant_config)

    def finalize(
        self,
        output: torch.Tensor,
        fused_expert_output: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        apply_router_weight_on_input: bool,
        apply_weights_and_reduce: bool = True
    ) -> None:
        pf = self._get_current_prepare_finalize()
        return pf.finalize(
            output, fused_expert_output, topk_weights, topk_ids,
            apply_router_weight_on_input, apply_weights_and_reduce)

    def finalize_async(
        self,
        output: torch.Tensor,
        fused_expert_output: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        apply_router_weight_on_input: bool,
        apply_weights_and_reduce: bool = True
    ):
        pf = self._get_current_prepare_finalize()
        if hasattr(pf, "finalize_async"):
            return pf.finalize_async(
                output, fused_expert_output, topk_weights, topk_ids,
                apply_router_weight_on_input, apply_weights_and_reduce)
        return pf.finalize(
            output, fused_expert_output, topk_weights, topk_ids,
            apply_router_weight_on_input, apply_weights_and_reduce)
