# Copyright (c) Alibaba, Inc. and its affiliates.
from dotenv import dotenv_values

env = dotenv_values('.env')

import unittest
from unittest import TestCase

from evalscope.config import TaskConfig
from evalscope.constants import EvalType, JudgeStrategy, OutputType
from evalscope.run import run_task
from evalscope.utils.logger import get_logger

logger = get_logger()


class TestReasoning(TestCase):
    """Benchmark evaluation test cases."""

    def setUp(self):
        """Setup common test configuration."""
        self.base_config = {
            'model': 'Qwen3-0.6B',
            'api_url': 'http://0.0.0.0:8801/v1',
            'api_key': env.get('DASHSCOPE_API_KEY'),
            'eval_type': EvalType.OPENAI_API,
            'eval_batch_size': 5,
            'limit': 5,
            'generation_config': {
                'max_tokens': 4096,
                'temperature': 0.0,
                'seed': 42,
                'parallel_tool_calls': True,
                'extra_body':{'chat_template_kwargs': {'enable_thinking': False}} # 关闭思考模式
            },
            'judge_strategy': JudgeStrategy.AUTO,
            'judge_model_args': {
                'model_id': 'qwen2.5-72b-instruct',
                'api_url': 'https://dashscope.aliyuncs.com/compatible-mode/v1',
                'api_key': env.get('DASHSCOPE_API_KEY'),
                'generation_config': {
                    'temperature': 0.0,
                    'max_tokens': 4096,
                }
            },
            'debug': True,
        }

    def _run_dataset_test(self, dataset_name, dataset_args=None, use_mock=False, **config_overrides):
        """Helper method to run test for a specific dataset."""
        config = self.base_config.copy()
        config['datasets'] = [dataset_name]

        if use_mock:
            config['eval_type'] = EvalType.MOCK_LLM

        # 应用配置覆盖
        config.update(config_overrides)

        if dataset_args:
            config['dataset_args'] = {dataset_name: dataset_args}

        task_cfg = TaskConfig(**config)
        run_task(task_cfg=task_cfg)

    def _run_dataset_load_test(self, dataset_name, dataset_args=None):
        """Helper method to test dataset loading."""

        self._run_dataset_test(dataset_name, dataset_args, use_mock=True, limit=None)

    # Math & Reasoning datasets
    def test_gsm8k(self):
        """Test GSM8K math reasoning dataset."""
        self._run_dataset_test('gsm8k')


if __name__ == '__main__':
    # Run specific test: python -m unittest test_eval.TestBenchmark.test_gsm8k
    # Run all tests: python -m unittest test_eval.TestBenchmark
    unittest.main()