# coding=utf-8
# Copyright 2025 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.

from collections.abc import Callable, Iterator
import contextlib
import dataclasses
import functools
import multiprocessing
import os
import typing
from unittest import mock

from etils import epath
import pytest
import tensorflow_datasets as tfds
from tensorflow_datasets.core import download
from tensorflow_datasets.core import utils
from tensorflow_datasets.scripts.cli import build as build_lib
from tensorflow_datasets.scripts.cli import main

type_utils = tfds.core.utils.type_utils

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


class DummyDatasetNoGenerate(tfds.testing.DummyDataset):

  def _generate_examples(self):
    if True:  # pylint: disable=using-constant-test
      raise NotImplementedError('Should not be called')
    yield

  @utils.classproperty
  @classmethod
  def url_infos(cls) -> dict[str, download.checksums.UrlInfo] | None:
    return {
        'http://data.org/file1.zip': download.checksums.UrlInfo(
            size=42,
            checksum='d45899d9a6a0e48afb250aac7ee3dc50e73e263687f15761d754515cd8284e0a',
            filename='file1.zip',
        ),
    }


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


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

  original_dl_and_prepare = tfds.core.DatasetBuilder.download_and_prepare

  # Unfortunately, `mock.Mock` removes `self` from `call_args`, so we have
  # to patch the function to record the generated_ds manually.
  # See:
  # https://stackoverflow.com/questions/64792295/how-to-get-self-instance-in-mock-mock-call-args
  queue = multiprocessing.Queue()

  def _download_and_prepare(self, *args, **kwargs):
    # Remove version from generated name (as only last version can be generated)
    full_name = '/'.join(self.info.full_name.split('/')[:-1])
    queue.put(full_name)
    if mock_download_and_prepare:
      return
    else:
      return original_dl_and_prepare(self, *args, **kwargs)

  with mock.patch.object(
      tfds.core.DatasetBuilder, 'download_and_prepare', _download_and_prepare
  ):
    main.main(args)
    queue.put(None)

  generated_ds_names = []
  while full_name := queue.get():
    generated_ds_names.append(full_name)

  return generated_ds_names


@pytest.fixture(name='build')
def mock_build(
    default_data_dir: epath.Path,
) -> Callable[[str, bool], list[str]]:
  """Returns a function to execute `tfds build`."""
  return functools.partial(_build, data_dir=default_data_dir)


def test_build_single(build):
  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.registered.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(build):
  # 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',
  ]


@pytest.mark.parametrize('num_processes', range(1, 4))
def test_build_parallel(build, num_processes):
  # Order is not guaranteed
  assert set(build(f'trivia_qa --num-processes={num_processes}')) == set([
      'trivia_qa/rc',
      'trivia_qa/rc.nocontext',
      'trivia_qa/unfiltered',
      'trivia_qa/unfiltered.nocontext',
  ])


def test_build_dataset_configs(build):
  # 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',
  ]

  # --config Json
  config_json = '{"name":"my_config","description":"abcd"}'
  assert build(f'imdb_reviews --config {config_json}') == [
      'imdb_reviews/my_config',
  ]

  # 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(build):
  # 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(build, default_data_dir: epath.Path):
  data_dir = 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_max_examples_per_split_0(build, default_data_dir: epath.Path):
  assert build(
      'dummy_dataset_no_generate --max_examples_per_split 0',
      mock_download_and_prepare=False,
  ) == ['dummy_dataset_no_generate']

  builder_path = default_data_dir / 'dummy_dataset_no_generate/1.0.0'
  # Dataset has been generated
  assert builder_path.exists()
  # tf-records files have not been generated
  assert sorted(p.name for p in builder_path.iterdir()) == [
      'dataset_info.json',
      'features.json',
  ]


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

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

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

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

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

  # cd .../datasets/ && tfds build dummy_dataset --imports=xxx
  # --imports is passed. so do not load dataset from file
  with mock_cwd(_DUMMY_DATASET_PATH.parent):
    with pytest.raises(tfds.core.registered.DatasetNotFoundError):
      assert build('dummy_dataset --imports=os')

  # cd .../datasets/ && tfds 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(build):
  # DummyDataset isn't registered by default
  with pytest.raises(tfds.core.registered.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']


def test_download_only(build):
  with mock.patch.object(
      download.DownloadManager, 'download', autospec=True
  ) as mock_download:
    assert not build('dummy_dataset_no_generate --download_only')
    mock_download.assert_called_with(
        mock.ANY, {'file0': 'http://data.org/file1.zip'}
    )


@pytest.mark.parametrize(
    'args,download_config_kwargs',
    [
        ('--manual_dir=/a/b', {'manual_dir': epath.Path('/a/b')}),
        (
            '--manual_dir=/a/b --add_name_to_manual_dir',
            {'manual_dir': epath.Path('/a/b/x')},
        ),
        ('--extract_dir=/a/b', {'extract_dir': epath.Path('/a/b')}),
        ('--max_examples_per_split=42', {'max_examples_per_split': 42}),
        ('--register_checksums', {'register_checksums': True}),
        ('--force_checksums_validation', {'force_checksums_validation': True}),
        ('--max_shard_size_mb=128', {'max_shard_size': 128 << 20}),
        (
            '--download_config={"max_shard_size":1234}',
            {'max_shard_size': 1234},
        ),
    ],
)
def test_make_download_config(args: str, download_config_kwargs):
  args = main._parse_flags(f'tfds build x {args}'.split())
  cmd_args = typing.cast(build_lib.Args, args.command)
  actual = build_lib._make_download_config(cmd_args, dataset_name='x')
  # Ignore the beam runner
  actual = actual.replace(beam_runner=None)
  expected = tfds.download.DownloadConfig(**download_config_kwargs)
  assert dataclasses.asdict(actual) == dataclasses.asdict(expected)
