#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#

"""A framework for developing sources for new file types.

To create a source for a new file type a sub-class of :class:`FileBasedSource`
should be created. Sub-classes of :class:`FileBasedSource` must implement the
method :meth:`FileBasedSource.read_records()`. Please read the documentation of
that method for more details.

For an example implementation of :class:`FileBasedSource` see
:class:`~apache_beam.io._AvroSource`.
"""

# pytype: skip-file

from typing import Callable
from typing import Iterable
from typing import Union

from apache_beam.internal import pickler
from apache_beam.io import concat_source
from apache_beam.io import iobase
from apache_beam.io import range_trackers
from apache_beam.io.filesystem import CompressionTypes
from apache_beam.io.filesystem import FileMetadata
from apache_beam.io.filesystems import FileSystems
from apache_beam.io.restriction_trackers import OffsetRange
from apache_beam.options.value_provider import StaticValueProvider
from apache_beam.options.value_provider import ValueProvider
from apache_beam.options.value_provider import check_accessible
from apache_beam.transforms.core import DoFn
from apache_beam.transforms.core import ParDo
from apache_beam.transforms.core import PTransform
from apache_beam.transforms.display import DisplayDataItem
from apache_beam.transforms.util import Reshuffle

MAX_NUM_THREADS_FOR_SIZE_ESTIMATION = 25

__all__ = ['FileBasedSource']


class FileBasedSource(iobase.BoundedSource):
  """A :class:`~apache_beam.io.iobase.BoundedSource` for reading a file glob of
  a given type."""

  MIN_NUMBER_OF_FILES_TO_STAT = 100
  MIN_FRACTION_OF_FILES_TO_STAT = 0.01

  def __init__(
      self,
      file_pattern,
      min_bundle_size=0,
      compression_type=CompressionTypes.AUTO,
      splittable=True,
      validate=True):
    """Initializes :class:`FileBasedSource`.

    Args:
      file_pattern (str): the file glob to read a string or a
        :class:`~apache_beam.options.value_provider.ValueProvider`
        (placeholder to inject a runtime value).
      min_bundle_size (int): minimum size of bundles that should be generated
        when performing initial splitting on this source.
      compression_type (str): Used to handle compressed output files.
        Typical value is :attr:`CompressionTypes.AUTO
        <apache_beam.io.filesystem.CompressionTypes.AUTO>`,
        in which case the final file path's extension will be used to detect
        the compression.
      splittable (bool): whether :class:`FileBasedSource` should try to
        logically split a single file into data ranges so that different parts
        of the same file can be read in parallel. If set to :data:`False`,
        :class:`FileBasedSource` will prevent both initial and dynamic splitting
        of sources for single files. File patterns that represent multiple files
        may still get split into sources for individual files. Even if set to
        :data:`True` by the user, :class:`FileBasedSource` may choose to not
        split the file, for example, for compressed files where currently it is
        not possible to efficiently read a data range without decompressing the
        whole file.
      validate (bool): Boolean flag to verify that the files exist during the
        pipeline creation time.

    Raises:
      TypeError: when **compression_type** is not valid or if
        **file_pattern** is not a :class:`str` or a
        :class:`~apache_beam.options.value_provider.ValueProvider`.
      ValueError: when compression and splittable files are
        specified.
      IOError: when the file pattern specified yields an empty
        result.
    """

    if not isinstance(file_pattern, (str, ValueProvider)):
      raise TypeError(
          '%s: file_pattern must be of type string'
          ' or ValueProvider; got %r instead' %
          (self.__class__.__name__, file_pattern))

    if isinstance(file_pattern, str):
      file_pattern = StaticValueProvider(str, file_pattern)
    self._pattern = file_pattern

    self._concat_source = None
    self._min_bundle_size = min_bundle_size
    if not CompressionTypes.is_valid_compression_type(compression_type):
      raise TypeError(
          'compression_type must be CompressionType object but '
          'was %s' % type(compression_type))
    self._compression_type = compression_type
    self._splittable = splittable
    if validate and file_pattern.is_accessible():
      self._validate()

  def display_data(self):
    return {
        'file_pattern': DisplayDataItem(
            str(self._pattern), label="File Pattern"),
        'compression': DisplayDataItem(
            str(self._compression_type), label='Compression Type')
    }

  @check_accessible(['_pattern'])
  def _get_concat_source(self) -> concat_source.ConcatSource:
    if self._concat_source is None:
      pattern = self._pattern.get()

      single_file_sources = []
      match_result = FileSystems.match([pattern])[0]
      files_metadata = match_result.metadata_list

      # We create a reference for FileBasedSource that will be serialized along
      # with each _SingleFileSource. To prevent this FileBasedSource from having
      # a reference to ConcatSource (resulting in quadratic space complexity)
      # we clone it here.
      file_based_source_ref = pickler.roundtrip(self)

      for file_metadata in files_metadata:
        file_name = file_metadata.path
        file_size = file_metadata.size_in_bytes
        if file_size == 0:
          continue  # Ignoring empty file.

        # We determine splittability of this specific file.
        splittable = (
            self.splittable and _determine_splittability_from_compression_type(
                file_name, self._compression_type))

        single_file_source = _SingleFileSource(
            file_based_source_ref,
            file_name,
            0,
            file_size,
            min_bundle_size=self._min_bundle_size,
            splittable=splittable)
        single_file_sources.append(single_file_source)

      FileSystems.report_source_lineage(pattern)
      self._concat_source = concat_source.ConcatSource(single_file_sources)

    return self._concat_source

  def open_file(self, file_name):
    return FileSystems.open(
        file_name,
        'application/octet-stream',
        compression_type=self._compression_type)

  @check_accessible(['_pattern'])
  def _validate(self):
    """Validate if there are actual files in the specified glob pattern
    """
    pattern = self._pattern.get()

    # Limit the responses as we only want to check if something exists
    match_result = FileSystems.match([pattern], limits=[1])[0]
    if len(match_result.metadata_list) <= 0:
      raise IOError('No files found based on the file pattern %s' % pattern)

  def split(
      self, desired_bundle_size=None, start_position=None, stop_position=None):
    return self._get_concat_source().split(
        desired_bundle_size=desired_bundle_size,
        start_position=start_position,
        stop_position=stop_position)

  def estimate_size(self):
    return self._get_concat_source().estimate_size()

  def read(self, range_tracker):
    return self._get_concat_source().read(range_tracker)

  def get_range_tracker(self, start_position, stop_position):
    return self._get_concat_source().get_range_tracker(
        start_position, stop_position)

  def read_records(self, file_name, offset_range_tracker):
    """Returns a generator of records created by reading file 'file_name'.

    Args:
      file_name: a ``string`` that gives the name of the file to be read. Method
                 ``FileBasedSource.open_file()`` must be used to open the file
                 and create a seekable file object.
      offset_range_tracker: a object of type ``OffsetRangeTracker``. This
                            defines the byte range of the file that should be
                            read. See documentation in
                            ``iobase.BoundedSource.read()`` for more information
                            on reading records while complying to the range
                            defined by a given ``RangeTracker``.

    Returns:
      an iterator that gives the records read from the given file.
    """
    raise NotImplementedError

  @property
  def splittable(self):
    return self._splittable


def _determine_splittability_from_compression_type(file_path, compression_type):
  if compression_type == CompressionTypes.AUTO:
    compression_type = CompressionTypes.detect_compression_type(file_path)

  return compression_type == CompressionTypes.UNCOMPRESSED


class _SingleFileSource(iobase.BoundedSource):
  """Denotes a source for a specific file type."""
  def __init__(
      self,
      file_based_source,
      file_name,
      start_offset,
      stop_offset,
      min_bundle_size=0,
      splittable=True):
    if not isinstance(start_offset, int):
      raise TypeError(
          'start_offset must be a number. Received: %r' % start_offset)
    if stop_offset != range_trackers.OffsetRangeTracker.OFFSET_INFINITY:
      if not isinstance(stop_offset, int):
        raise TypeError(
            'stop_offset must be a number. Received: %r' % stop_offset)
      if start_offset >= stop_offset:
        raise ValueError(
            'start_offset must be smaller than stop_offset. Received %d and %d '
            'for start and stop offsets respectively' %
            (start_offset, stop_offset))

    self._file_name = file_name
    self._is_gcs_file = file_name.startswith('gs://') if file_name else False
    self._start_offset = start_offset
    self._stop_offset = stop_offset
    self._min_bundle_size = min_bundle_size
    self._file_based_source = file_based_source
    self._splittable = splittable

  def split(self, desired_bundle_size, start_offset=None, stop_offset=None):
    if start_offset is None:
      start_offset = self._start_offset
    if stop_offset is None:
      stop_offset = self._stop_offset

    if self._splittable:
      splits = OffsetRange(start_offset, stop_offset).split(
          desired_bundle_size, self._min_bundle_size)
      for split in splits:
        yield iobase.SourceBundle(
            split.stop - split.start,
            _SingleFileSource(
                # Copying this so that each sub-source gets a fresh instance.
                pickler.roundtrip(self._file_based_source),
                self._file_name,
                split.start,
                split.stop,
                min_bundle_size=self._min_bundle_size,
                splittable=self._splittable),
            split.start,
            split.stop)
    else:
      # Returning a single sub-source with end offset set to OFFSET_INFINITY (so
      # that all data of the source gets read) since this source is
      # unsplittable. Choosing size of the file as end offset will be wrong for
      # certain unsplittable source, e.g., compressed sources.
      yield iobase.SourceBundle(
          stop_offset - start_offset,
          _SingleFileSource(
              self._file_based_source,
              self._file_name,
              start_offset,
              range_trackers.OffsetRangeTracker.OFFSET_INFINITY,
              min_bundle_size=self._min_bundle_size,
              splittable=self._splittable),
          start_offset,
          range_trackers.OffsetRangeTracker.OFFSET_INFINITY)

  def estimate_size(self):
    return self._stop_offset - self._start_offset

  def get_range_tracker(self, start_position, stop_position):
    if start_position is None:
      start_position = self._start_offset
    if stop_position is None:
      # If file is unsplittable we choose OFFSET_INFINITY as the default end
      # offset so that all data of the source gets read. Choosing size of the
      # file as end offset will be wrong for certain unsplittable source, for
      # e.g., compressed sources.
      stop_position = (
          self._stop_offset if self._splittable else
          range_trackers.OffsetRangeTracker.OFFSET_INFINITY)

    range_tracker = range_trackers.OffsetRangeTracker(
        start_position, stop_position)
    if not self._splittable:
      range_tracker = range_trackers.UnsplittableRangeTracker(range_tracker)

    return range_tracker

  def read(self, range_tracker):
    return self._file_based_source.read_records(self._file_name, range_tracker)

  def default_output_coder(self):
    return self._file_based_source.default_output_coder()


class _ExpandIntoRanges(DoFn):
  def __init__(
      self, splittable, compression_type, desired_bundle_size, min_bundle_size):
    self._desired_bundle_size = desired_bundle_size
    self._min_bundle_size = min_bundle_size
    self._splittable = splittable
    self._compression_type = compression_type
    self._size_track = None

  def process(self, element: Union[str, FileMetadata], *args,
              **kwargs) -> Iterable[tuple[FileMetadata, OffsetRange]]:
    if isinstance(element, FileMetadata):
      metadata_list = [element]
    else:
      match_results = FileSystems.match([element])
      metadata_list = match_results[0].metadata_list
    for metadata in metadata_list:
      FileSystems.report_source_lineage(metadata.path)

      splittable = (
          self._splittable and _determine_splittability_from_compression_type(
              metadata.path, self._compression_type))

      if splittable:
        for split in OffsetRange(0, metadata.size_in_bytes).split(
            self._desired_bundle_size, self._min_bundle_size):
          yield (metadata, split)
      else:
        yield (
            metadata,
            OffsetRange(0, range_trackers.OffsetRangeTracker.OFFSET_INFINITY))


class _ReadRange(DoFn):
  def __init__(
      self,
      source_from_file: Union[str, iobase.BoundedSource],
      with_filename: bool = False) -> None:
    self._source_from_file = source_from_file
    self._with_filename = with_filename

  def process(self, element, *args, **kwargs):
    metadata, range = element
    source = self._source_from_file(metadata.path)
    # Following split() operation has to be performed to create a proper
    # _SingleFileSource. Otherwise what we have is a ConcatSource that contains
    # a single _SingleFileSource. ConcatSource.read() expects a RangeTracker for
    # sub-source range and reads full sub-sources (not byte ranges).
    source_list = list(source.split(float('inf')))
    # Handle the case of an empty source.
    if not source_list:
      return
    source = source_list[0].source

    for record in source.read(range.new_tracker()):
      if self._with_filename:
        yield (metadata.path, record)
      else:
        yield record


class ReadAllFiles(PTransform):
  """A Read transform that reads a PCollection of files.

  Pipeline authors should not use this directly. This is to be used by Read
  PTransform authors who wishes to implement file-based Read transforms that
  read a PCollection of files.
  """
  def __init__(
      self,
      splittable: bool,
      compression_type,
      desired_bundle_size: int,
      min_bundle_size: int,
      source_from_file: Callable[[str], iobase.BoundedSource],
      with_filename: bool = False):
    """
    Args:
      splittable: If False, files won't be split into sub-ranges. If True,
                  files may or may not be split into data ranges.
      compression_type: A ``CompressionType`` object that specifies the
                  compression type of the files that will be processed. If
                  ``CompressionType.AUTO``, system will try to automatically
                  determine the compression type based on the extension of
                  files.
      desired_bundle_size: the desired size of data ranges that should be
                           generated when splitting a file into data ranges.
      min_bundle_size: minimum size of data ranges that should be generated when
                           splitting a file into data ranges.
      source_from_file: a function that produces a ``BoundedSource`` given a
                        file name. System will use this function to generate
                        ``BoundedSource`` objects for file paths. Note that file
                        paths passed to this will be for individual files, not
                        for file patterns even if the ``PCollection`` of files
                        processed by the transform consist of file patterns.
      with_filename: If True, returns a Key Value with the key being the file
        name and the value being the actual data. If False, it only returns
        the data.
    """
    self._splittable = splittable
    self._compression_type = compression_type
    self._desired_bundle_size = desired_bundle_size
    self._min_bundle_size = min_bundle_size
    self._source_from_file = source_from_file
    self._with_filename = with_filename
    # TODO(BEAM-14497) always reshuffle once gbk always trigger works.
    self._is_reshuffle = True

  def _disable_reshuffle(self):
    # TODO(BEAM-14497) Remove this private method once gbk always trigger works.
    #
    # Currently Reshuffle() holds elements until the stage is completed. When
    # ReadRange is needed instantly after match (like read continuously), the
    # reshard is temporarily disabled. However, the read then does not scale and
    # is deemed experimental.
    self._is_reshuffle = False
    return self

  def expand(self, pvalue):
    pvalue = (
        pvalue
        | 'ExpandIntoRanges' >> ParDo(
            _ExpandIntoRanges(
                self._splittable,
                self._compression_type,
                self._desired_bundle_size,
                self._min_bundle_size)))
    if self._is_reshuffle:
      pvalue = pvalue | 'Reshard' >> Reshuffle()
    return (
        pvalue
        | 'ReadRange' >> ParDo(
            _ReadRange(
                self._source_from_file, with_filename=self._with_filename)))