# coding=utf-8
# Copyright 2020 The TensorFlow Datasets Authors.
#
# 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.

"""Tests for tensorflow_datasets.scripts.cli.build."""

import contextlib
import os
import pathlib
from typing import Iterator, List
from unittest import mock

import pytest

import tensorflow_datasets as tfds
from tensorflow_datasets.scripts.cli import main

type_utils = tfds.core.utils.type_utils

_DUMMY_DATASET_PATH = tfds.core.tfds_path() / 'testing/dummy_dataset'


@pytest.fixture(scope='function', autouse=True)
def mock_default_data_dir(tmp_path: pathlib.Path):
  """Changes the default `--data_dir` to tmp_path."""
  tmp_path = tmp_path / 'datasets'
  default_data_dir = os.environ.get('TFDS_DATA_DIR')
  try:
    os.environ['TFDS_DATA_DIR'] = os.fspath(tmp_path)
    yield tmp_path
  finally:
    if default_data_dir:
      os.environ['TFDS_DATA_DIR'] = default_data_dir
    else:
      del os.environ['TFDS_DATA_DIR']


@contextlib.contextmanager
def mock_cwd(path: type_utils.PathLike) -> Iterator[None]:
  """Mock the current directory."""
  path = pathlib.Path(path)
  assert path.exists() and path.is_dir()  # Check given path is valid cwd dir
  with mock.patch('os.getcwd', return_value=os.fspath(path)):
    yield


def _build(cmd_flags: str) -> List[str]:
  """Executes `tfds build {cmd_flags}` and returns the list of generated ds."""
  # Execute the command
  args = main._parse_flags(f'tfds build {cmd_flags}'.split())

  # Unfortunatelly, `mock.Mock` remove `self` from `call_args`, so we have
  # to patch manually the function to record the generated_ds.
  # See:
  # https://stackoverflow.com/questions/64792295/how-to-get-self-instance-in-mock-mock-call-args
  generated_ds_names = []
  def _download_and_prepare(self, *args, **kwargs):
    del args, kwargs
    # Remove version from generated name (as only last version can be generated)
    full_name = '/'.join(self.info.full_name.split('/')[:-1])
    generated_ds_names.append(full_name)
    return

  with mock.patch(
      'tensorflow_datasets.core.DatasetBuilder.download_and_prepare',
      _download_and_prepare,
  ):
    main.main(args)
  return generated_ds_names


def test_build_single():
  assert _build('mnist') == ['mnist']
  assert _build('mnist:3.0.1') == ['mnist']
  # Keyword arguments also possible
  assert _build('--datasets mnist') == ['mnist']

  with pytest.raises(tfds.core.load.DatasetNotFoundError):
    _build('unknown_dataset')

  with pytest.raises(AssertionError, match='cannot be loaded at version 1.0.0'):
    _build('mnist:1.0.0')  # Can only built the last version

  with pytest.raises(ValueError, match='not have config'):
    _build('mnist --config_idx 0')


def test_build_multiple():
  # Multiple datasets can be built in a single call
  assert _build('mnist imagenet2012 cifar10') == [
      'mnist',
      'imagenet2012',
      'cifar10',
  ]
  # Keyword arguments also possible
  assert _build('mnist --datasets imagenet2012 cifar10') == [
      'mnist',
      'imagenet2012',
      'cifar10',
  ]


def test_build_dataset_configs():
  # By default, all configs are build
  assert _build('trivia_qa') == [
      'trivia_qa/rc',
      'trivia_qa/rc.nocontext',
      'trivia_qa/unfiltered',
      'trivia_qa/unfiltered.nocontext',
  ]

  # If config is set, only the defined config is generated

  # --config_idx
  assert _build('trivia_qa --config_idx=0') == ['trivia_qa/rc']

  # --config
  assert _build('trivia_qa --config unfiltered.nocontext') == [
      'trivia_qa/unfiltered.nocontext',
  ]

  # name/config
  assert _build('trivia_qa/unfiltered.nocontext') == [
      'trivia_qa/unfiltered.nocontext'
  ]

  with pytest.raises(ValueError, match='Config should only be defined once'):
    _build('trivia_qa/unfiltered.nocontext --config_idx=0')

  with pytest.raises(ValueError, match='greater than number of configs'):
    _build('trivia_qa --config_idx 100')


def test_exclude_datasets():
  # Exclude all datasets except 2
  all_ds = [b for b in tfds.list_builders() if b not in ('mnist', 'cifar10')]
  all_ds_str = ','.join(all_ds)

  assert _build(f'--exclude_datasets {all_ds_str}') == [
      'cifar10',
      'mnist',
  ]

  with pytest.raises(ValueError, match='--exclude_datasets can\'t be used'):
    _build('mnist --exclude_datasets cifar10')


def test_build_overwrite(mock_default_data_dir: pathlib.Path):  # pylint: disable=redefined-outer-name
  data_dir = mock_default_data_dir / 'mnist/3.0.1'
  data_dir.mkdir(parents=True)
  metadata_path = tfds.core.tfds_path(
      'testing/test_data/dataset_info/mnist/3.0.1'
  )

  for f in metadata_path.iterdir():  # Copy metadata files.
    data_dir.joinpath(f.name).write_text(f.read_text())

  # By default, will skip generation if the data already exists
  assert _build('mnist') == ['mnist']  # Called, but no-op
  assert data_dir.exists()

  assert _build('mnist --overwrite') == ['mnist']
  assert not data_dir.exists()  # Previous data-dir has been removed


def test_build_files():
  # Make sure DummyDataset isn't registered by default
  with pytest.raises(tfds.core.load.DatasetNotFoundError):
    _build('dummy_dataset')

  with pytest.raises(FileNotFoundError, match='Could not find .* script'):
    _build('')

  # cd .../datasets/dummy_dataset && gtfds build
  with mock_cwd(_DUMMY_DATASET_PATH):
    assert _build('') == ['dummy_dataset']

  # cd .../datasets/dummy_dataset && gtfds build dummy_dataset.py
  with mock_cwd(_DUMMY_DATASET_PATH):
    assert _build('dummy_dataset.py') == ['dummy_dataset']

  # cd .../datasets/ && gtfds build dummy_dataset
  with mock_cwd(_DUMMY_DATASET_PATH.parent):
    assert _build('dummy_dataset') == ['dummy_dataset']

  # cd .../datasets/ && gtfds build dummy_dataset/dummy_dataset
  with mock_cwd(_DUMMY_DATASET_PATH.parent):
    assert _build('dummy_dataset/dummy_dataset') == ['dummy_dataset']


# Somehow only with tf-nightly, `dummy_dataset` is already imported by
# community/load_test.py (with `skip_registration()`). Thus import here have
# no-effects.
@pytest.mark.skip(reason='Conflict with `load_test.py`')
def test_build_import():
  # DummyDataset isn't registered by default
  with pytest.raises(tfds.core.load.DatasetNotFoundError):
    _build('dummy_dataset')

  # --imports register the dataset
  ds_module = 'tensorflow_datasets.testing.dummy_dataset.dummy_dataset'
  assert _build(f'dummy_dataset --imports {ds_module}') == ['dummy_dataset']
