# Copyright 2025 Huawei Technologies Co., Ltd
#
# 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.
# ============================================================================
"""Test Mcore Qwen3Moe inference"""
import os
import random
from pathlib import Path

import pytest

from mindformers.tools.logger import logger
from tests.st.test_multi_cards_cases.utils import TaskType

_LEVEL_0_TASK_TIME = 150
_LEVEL_1_TASK_TIME = 0
_TASK_TYPE = TaskType.FOUR_CARDS_TASK


class TestQuantDeepseekParallelInference:
    """Test class for quant deepseek in inference"""

    def setup_method(self):
        """Setup method to prepare test environment"""
        self.sh_path = Path(__file__).parent.resolve()
        self.run_script_path = self.sh_path / "run_quant_deepseek_infer.py"
        assert self.run_script_path.exists(), f"Run script not found: {self.run_script_path}"

    @pytest.mark.level0
    @pytest.mark.parametrize("quant_algo", ['A8W8'])
    def test_quant_deepseek_level0(self, quant_algo):
        """Test two cards for quant deepseek."""
        port_id = int(os.environ.get("ASCEND_PORT_ID", random.randint(50000, 65535)))
        cmd_list = [
            "msrun",
            f"--worker_num=4",
            f"--local_worker_num=4",  # Should match NPU cards available
            f"--master_port={port_id}",  # Ensure port is unique per test run if parallelized at pytest level
            f"--log_dir=./msrun_log_{quant_algo}_quant_deepseek_infer",
            "--join=True"]
        cmd_list += [
            str(self.run_script_path),
            f"--quant_algo={quant_algo}"
        ]
        cmd = " ".join(cmd_list)
        logger.info(f"Running command: {cmd}")
        return_code = os.system(cmd)
        assert return_code == 0, "deepseek inference st failed."
