#
# Copyright (c) 2025 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.
#

# Patch vllm's FusedMoE factory to use AscendMoERunner by default.
#
# vllm's FusedMoE is a factory function (not a class). deepseek_v2 and other
# models do `from vllm.model_executor.layers.fused_moe import FusedMoE` and
# call it directly, so we must patch the binding in the package __init__ as
# well as the layer module before any model is imported.
#
# Import order in worker.__init__:
#   1. adapt_patch()  ->  this file runs  ->  FusedMoE patched
#   2. from vllm_ascend import ops
#   3. model loading  ->  deepseek_v2 imported  ->  gets patched FusedMoE  ✓

from vllm_ascend.utils import is_310p, vllm_version_is

if not vllm_version_is("0.23.0"):
    import vllm.model_executor.layers.fused_moe as _fused_moe_pkg
    import vllm.model_executor.layers.fused_moe.layer as _fused_moe_layer

    # Capture the real original before fused_moe.py's module-level code runs.
    _original_FusedMoE = _fused_moe_layer.FusedMoE

    if is_310p():
        from vllm_ascend._310p.fused_moe.fused_moe import AscendMoERunner310 as _DefaultAscendMoERunner
    else:
        from vllm_ascend.ops.fused_moe.fused_moe import AscendMoERunner as _DefaultAscendMoERunner

    def _ascend_FusedMoE(*args, runner_cls=None, runner_args=None, **kwargs):
        if runner_cls is None:
            runner_cls = _DefaultAscendMoERunner
        # 'hash' is a DeepSeek V4 flag already consumed before FusedMoE is called;
        # 'tid2eid' is Ascend-specific and must reach AscendMoERunner via runner_args.
        kwargs.pop("hash", None)
        tid2eid = kwargs.pop("tid2eid", None)
        if tid2eid is not None:
            runner_args = dict(runner_args) if runner_args is not None else {}
            runner_args["tid2eid"] = tid2eid
        return _original_FusedMoE(*args, runner_cls=runner_cls, runner_args=runner_args, **kwargs)

    _fused_moe_layer.FusedMoE = _ascend_FusedMoE
    _fused_moe_pkg.FusedMoE = _ascend_FusedMoE