# Copyright 2021 Zilliz. All rights reserved.
#
# 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.
import os
from setuptools import setup, find_packages
from setuptools.command.install import install
import unittest
from pathlib import Path
import shutil


def test_suite():
    test_loader = unittest.TestLoader()
    test_suite = test_loader.discover('towhee/tests', pattern='test_*.py')
    return test_suite


def create_towhee_cache(src: str, dst: str, is_dir: bool = False):
    if not Path(dst).is_dir():
        Path(dst).mkdir(parents=True)
        if is_dir:
            dst = os.path.join(dst, src.split('/')[-1])
            shutil.copytree(src, dst)
        else:
            shutil.copy(src, dst)


class PostInstallCommand(install):
    def run(self):
        from towhee.engine import LOCAL_OPERATOR_CACHE, LOCAL_PIPELINE_CACHE, MOCK_OPS, MOCK_PIPELINES
        install.run(self)
        create_towhee_cache(MOCK_PIPELINES, LOCAL_PIPELINE_CACHE)
        create_towhee_cache(MOCK_OPS, LOCAL_OPERATOR_CACHE, True)


setup(name="towhee_test",
      version="0.1.0",
      description="",
      author="Towhee Team",
      author_email="towhee-team@zilliz.com",
      url="https://github.com/towhee-io/towhee",
      test_suite="setup.test_suite",
      cmdclass={'install': PostInstallCommand},
      install_requires=[
          'torch>=1.2.0',
          'torchvision>=0.4.0',
          'numpy>=1.20.3',
          'pandas>=1.2.4',
          'pyyaml>=5.3',
          'requests>=2.25',
          'tqdm>=4.59.0',
          'pillow>=8.3.1',
          'scipy>=1.7.1',
      ],
      packages=find_packages(),
      package_data={'towhee.tests.test_util': ['*.yaml']},
      license="http://www.apache.org/licenses/LICENSE-2.0")
