# Copyright (c) Alibaba, Inc. and its affiliates.
import os
import pytest
import subprocess
import unittest

pytestmark = pytest.mark.timeout(600)

from evalscope.run import run_task
from evalscope.utils.import_utils import is_module_installed
from evalscope.utils.logger import get_logger
from tests.utils import test_level_list

logger = get_logger()


class TestCLIPBenchmark(unittest.TestCase):

    def setUp(self) -> None:
        self._check_env('webdataset')

    def tearDown(self) -> None:
        pass

    @staticmethod
    def _check_env(module_name: str):
        if is_module_installed(module_name):
            logger.info(f'{module_name} is installed.')
        else:
            raise ModuleNotFoundError(f'run: pip install {module_name}')

    @unittest.skipUnless(0 in test_level_list(), 'skip test in current test level')
    def test_run_task(self):
        task_cfg = {
            'eval_backend': 'RAGEval',
            'eval_config': {
                'tool': 'clip_benchmark',
                'eval': {
                    'models': [
                        {
                            'model_name': 'AI-ModelScope/chinese-clip-vit-large-patch14-336px',
                        }
                    ],
                    'dataset_name': [
                        'muge',
                        'mnist',
                        'flickr8k'
                    ],
                    'split': 'test',
                    'batch_size': 128,
                    'num_workers': 1,
                    'verbose': True,
                    'skip_existing': False,
                    'cache_dir': 'cache',
                    'limit': 10,
                },
            },
        }

        run_task(task_cfg)

    @unittest.skipUnless(0 in test_level_list(), 'skip test in current test level')
    def test_run_custom(self):
        task_cfg = {
            'eval_backend': 'RAGEval',
            'eval_config': {
                'tool': 'clip_benchmark',
                'eval': {
                    'models': [
                        {
                            'model_name': 'AI-ModelScope/chinese-clip-vit-large-patch14-336px',
                        }
                    ],
                    'dataset_name': ['custom'],
                    'data_dir': 'custom_eval/multimodal/text-image-retrieval',
                    'split': 'test',
                    'batch_size': 128,
                    'num_workers': 1,
                    'verbose': True,
                    'skip_existing': False,
                    'limit': 10,
                },
            },
        }

        run_task(task_cfg)


if __name__ == '__main__':
    unittest.main(buffer=False)
