# coding=utf-8
# Copyright 2023 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 file_utils."""

import os
import time
from unittest import mock

from absl.testing import flagsaver
from etils import epath
import pytest
from tensorflow_datasets import testing
from tensorflow_datasets.core import naming
from tensorflow_datasets.core.utils import file_utils
from tensorflow_datasets.core.utils import read_config


def test_default_data_dir():
  data_dir = file_utils.get_default_data_dir(given_data_dir=None)
  assert data_dir
  assert isinstance(data_dir, str)


def test_list_dataset_variants_with_configs(mock_fs: testing.MockFs):
  data_dir = epath.Path('/a')
  dataset_dir = data_dir / 'ds'
  configs_and_versions = {
      'x': ['1.0.0', '1.0.1'],
      'y': ['2.0.0'],
  }
  for config, versions in configs_and_versions.items():
    for version in versions:
      mock_fs.add_file(dataset_dir / config / version / 'dataset_info.json')
      mock_fs.add_file(dataset_dir / config / version / 'features.json')

  references = sorted(
      file_utils.list_dataset_variants(
          dataset_name='my_ds', dataset_dir=dataset_dir
      )
  )
  assert references == [
      naming.DatasetReference(
          dataset_name='my_ds', config='x', version='1.0.0', data_dir=data_dir
      ),
      naming.DatasetReference(
          dataset_name='my_ds', config='x', version='1.0.1', data_dir=data_dir
      ),
      naming.DatasetReference(
          dataset_name='my_ds', config='y', version='2.0.0', data_dir=data_dir
      ),
  ]


def test_list_dataset_variants_with_configs_no_versions(
    mock_fs: testing.MockFs,
):
  data_dir = epath.Path('/a')
  dataset_dir = data_dir / 'ds'
  configs_and_versions = {
      'x': ['1.0.0', '1.0.1'],
      'y': ['2.0.0'],
  }
  for config, versions in configs_and_versions.items():
    for version in versions:
      mock_fs.add_file(dataset_dir / config / version / 'dataset_info.json')
      mock_fs.add_file(dataset_dir / config / version / 'features.json')

  references = sorted(
      file_utils.list_dataset_variants(
          dataset_name='my_ds', dataset_dir=dataset_dir, include_versions=False
      )
  )
  assert references == [
      naming.DatasetReference(
          dataset_name='my_ds', config='x', data_dir=data_dir
      ),
      naming.DatasetReference(
          dataset_name='my_ds', config='y', data_dir=data_dir
      ),
  ]


def test_list_dataset_variants_without_configs(mock_fs: testing.MockFs):
  data_dir = epath.Path('/a')
  dataset_dir = data_dir / 'ds'
  # Version 1.0.0 doesn't have features.json, because it was generated with an
  # old version of TFDS.
  mock_fs.add_file(dataset_dir / '1.0.0' / 'dataset_info.json')
  mock_fs.add_file(dataset_dir / '1.0.1' / 'dataset_info.json')
  mock_fs.add_file(dataset_dir / '1.0.1' / 'features.json')

  # List dirs including datasets generated by old TFDS versions.
  references = sorted(
      file_utils.list_dataset_variants(
          dataset_name='my_ds',
          dataset_dir=dataset_dir,
          include_versions=True,
          include_old_tfds_version=True,
      )
  )
  assert references == [
      naming.DatasetReference(
          dataset_name='my_ds', version='1.0.0', data_dir=data_dir
      ),
      naming.DatasetReference(
          dataset_name='my_ds', version='1.0.1', data_dir=data_dir
      ),
  ]

  # List dirs excluding datasets generated by old TFDS versions.
  references = sorted(
      file_utils.list_dataset_variants(
          dataset_name='my_ds',
          dataset_dir=dataset_dir,
          include_versions=True,
          include_old_tfds_version=False,
      )
  )
  assert references == [
      naming.DatasetReference(
          dataset_name='my_ds', version='1.0.1', data_dir=data_dir
      )
  ]


def test_list_datasets_in_data_dir(mock_fs: testing.MockFs):
  data_dir = epath.Path('/a')
  mock_fs.add_file(data_dir / 'ds1/config1/1.0.0/dataset_info.json')
  mock_fs.add_file(data_dir / 'ds1/config1/1.0.0/features.json')
  mock_fs.add_file(data_dir / 'ds1/config1/2.0.0/dataset_info.json')
  mock_fs.add_file(data_dir / 'ds1/config1/2.0.0/features.json')
  mock_fs.add_file(data_dir / 'ds1/config2/1.0.0/dataset_info.json')
  mock_fs.add_file(data_dir / 'ds1/config2/1.0.0/features.json')
  mock_fs.add_file(data_dir / 'ds2/1.0.0/dataset_info.json')
  mock_fs.add_file(data_dir / 'ds2/1.0.0/features.json')

  # The following are problematic and should thus be ignored.
  mock_fs.add_file(
      os.path.join(data_dir, 'invalid-name/1.0.0/features.json'), content='x'
  )
  mock_fs.add_file(
      os.path.join(data_dir, 'invalid_version1/1.a.b/features.json'),
      content='x',
  )
  mock_fs.add_file(
      os.path.join(data_dir, 'invalid_version2/1.2.3.4/features.json'),
      content='x',
  )

  references = sorted(
      file_utils.list_datasets_in_data_dir(data_dir=epath.Path(data_dir))
  )
  data_dir = epath.Path('/a')
  assert references == [
      naming.DatasetReference(
          dataset_name='ds1',
          config='config1',
          version='1.0.0',
          data_dir=data_dir,
      ),
      naming.DatasetReference(
          dataset_name='ds1',
          config='config1',
          version='2.0.0',
          data_dir=data_dir,
      ),
      naming.DatasetReference(
          dataset_name='ds1',
          config='config2',
          version='1.0.0',
          data_dir=data_dir,
      ),
      naming.DatasetReference(
          dataset_name='ds2', version='1.0.0', data_dir=data_dir
      ),
  ]


def test_list_datasets_in_data_dir_with_namespace(mock_fs: testing.MockFs):
  namespace = 'ns'
  data_dir = epath.Path('/a')
  mock_fs.add_file(data_dir / 'ds1/config1/1.0.0/dataset_info.json')
  mock_fs.add_file(data_dir / 'ds1/config1/1.0.0/features.json')

  references = sorted(
      file_utils.list_datasets_in_data_dir(
          data_dir=epath.Path(data_dir),
          namespace=namespace,
          include_configs=True,
          include_versions=True,
      )
  )
  data_dir = epath.Path('/a')
  assert references == [
      naming.DatasetReference(
          dataset_name='ds1',
          namespace=namespace,
          config='config1',
          version='1.0.0',
          data_dir=data_dir,
      ),
  ]


@pytest.mark.parametrize(
    ['filename', 'result'],
    [
        ('abc', False),
        ('dataset_info.json', True),
        ('features.json', True),
        ('mnist-test.tfrecord-00000-of-00001', True),
        ('mnist-test.arrayrecord-00000-of-00001', True),
    ],
)
def METHOD_NAME(filename, result):
  assert file_utils._looks_like_a_tfds_file(filename) == result


@pytest.mark.parametrize(
    ['path', 'glob_result', 'expected'],
    [
        ('/a/*', ['/a/b', '/a/c'], ['/a/b', '/a/c']),
        ('/a/b', None, ['/a/b']),
        ('a/*', None, ['a/*']),
        ('/a/b@*', None, ['/a/b@*']),
    ],
)
def test_expand_glob(path, glob_result, expected):
  with mock.patch.object(epath, 'Path') as mock_epath:
    mock_epath.return_value.expanduser.return_value = path
    mock_epath.return_value.glob.return_value = glob_result
    actual = file_utils.expand_glob(path)
    if glob_result is not None:
      mock_epath.return_value.glob.assert_called_once_with(path[1:])
    else:
      mock_epath.return_value.glob.assert_not_called()
    actual = [os.fspath(p) for p in actual]
    assert actual == expected

if __name__ == '__main__':
  testing.test_main()