#
# 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.
#

"""``PTransforms`` for manipulating files in Apache Beam.

Provides reading ``PTransform``\\s, ``MatchFiles``,
``MatchAll``, that produces a ``PCollection`` of records representing a file
and its metadata; and ``ReadMatches``, which takes in a ``PCollection`` of file
metadata records, and produces a ``PCollection`` of ``ReadableFile`` objects.
These transforms currently do not support splitting by themselves.

Writing to Files
================

The transforms in this file include ``WriteToFiles``, which allows you to write
a ``beam.PCollection`` to files, and gives you many options to customize how to
do this.

The ``WriteToFiles`` transform supports bounded and unbounded PCollections
(i.e. it can be used both batch and streaming pipelines). For streaming
pipelines, it currently does not have support for multiple trigger firings
on the same window.

File Naming
-----------
One of the parameters received by ``WriteToFiles`` is a function specifying how
to name the files that are written. This is a function that takes in the
following parameters:

- window
- pane
- shard_index
- total_shards
- compression
- destination

It should return a file name that is unique for a combination of these
parameters.

The default naming strategy is to name files
in the format
`$prefix-$start-$end-$pane-$shard-of-$numShards$suffix$compressionSuffix`,
where:

- `$prefix` is, by default, `"output"`.
- `$start` and `$end` are the boundaries of the window for the data being
  written. These are omitted if we're using the Global window.
- `$pane` is the index for the number of firing for a window.
- `$shard` and `$numShards` are the current shard number, and the total number
  of shards for this window firing.
- `$suffix` is, by default, an empty string, but it can be set by the user via
  ``default_file_naming``.

Dynamic Destinations
--------------------
If the elements in the input ``beam.PCollection`` can be partitioned into groups
that should be treated differently (e.g. some events are to be stored as CSV,
while some others are to be stored as Avro files), it is possible to do this
by passing a `destination` parameter to ``WriteToFiles``. Something like the
following::

    my_pcollection | beam.io.fileio.WriteToFiles(
          path='/my/file/path',
          destination=lambda record: 'avro' if record['type'] == 'A' else 'csv',
          sink=lambda dest: AvroSink() if dest == 'avro' else CsvSink(),
          file_naming=beam.io.fileio.destination_prefix_naming())

In this transform, depending on the type of a record, it will be written down to
a destination named `'avro'`, or `'csv'`. The value returned by the
`destination` call is then passed to the `sink` call, to determine what sort of
sink will be used for each destination. The return type of the `destination`
parameter can be anything, as long as elements can be grouped by it.
"""

# pytype: skip-file

import collections
import logging
import random
import uuid
from collections import namedtuple
from typing import Any
from typing import BinaryIO  # pylint: disable=unused-import
from typing import Callable
from typing import Iterable
from typing import Optional
from typing import Union

import apache_beam as beam
from apache_beam.coders.coders import VarIntCoder
from apache_beam.io import filesystem
from apache_beam.io import filesystems
from apache_beam.io.filesystem import BeamIOError
from apache_beam.io.filesystem import CompressionTypes
from apache_beam.io.watch import PollFn
from apache_beam.io.watch import PollResult
from apache_beam.io.watch import TerminationCondition
from apache_beam.io.watch import Watch
from apache_beam.io.watch import never
from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.value_provider import StaticValueProvider
from apache_beam.options.value_provider import ValueProvider
from apache_beam.transforms.periodicsequence import PeriodicImpulse
from apache_beam.transforms.window import BoundedWindow
from apache_beam.transforms.window import FixedWindows
from apache_beam.transforms.window import GlobalWindow
from apache_beam.transforms.window import IntervalWindow
from apache_beam.transforms.window import TimestampedValue
from apache_beam.utils.timestamp import MAX_TIMESTAMP
from apache_beam.utils.timestamp import Duration
from apache_beam.utils.timestamp import Timestamp

__all__ = [
    'EmptyMatchTreatment',
    'MatchFiles',
    'MatchAll',
    'MatchContinuously',
    'ReadableFile',
    'ReadMatches',
    'WriteToFiles'
]

_LOGGER = logging.getLogger(__name__)

FileMetadata = namedtuple("FileMetadata", "mime_type compression_type")

CreateFileMetadataFn = Callable[[str, str], FileMetadata]


class EmptyMatchTreatment(object):
  """How to treat empty matches in ``MatchAll`` and ``MatchFiles`` transforms.

  If empty matches are disallowed, an error will be thrown if a pattern does not
  match any files."""

  ALLOW = 'ALLOW'
  DISALLOW = 'DISALLOW'
  ALLOW_IF_WILDCARD = 'ALLOW_IF_WILDCARD'

  @staticmethod
  def allow_empty_match(pattern, setting):
    if setting == EmptyMatchTreatment.ALLOW:
      return True
    elif setting == EmptyMatchTreatment.ALLOW_IF_WILDCARD and '*' in pattern:
      return True
    elif setting == EmptyMatchTreatment.DISALLOW:
      return False
    else:
      raise ValueError(setting)


class _MatchAllFn(beam.DoFn):
  def __init__(self, empty_match_treatment):
    self._empty_match_treatment = empty_match_treatment

  def process(self, file_pattern: str) -> list[filesystem.FileMetadata]:
    # TODO: Should we batch the lookups?
    match_results = filesystems.FileSystems.match([file_pattern])
    match_result = match_results[0]

    if (not match_result.metadata_list and
        not EmptyMatchTreatment.allow_empty_match(file_pattern,
                                                  self._empty_match_treatment)):
      raise BeamIOError(
          'Empty match for pattern %s. Disallowed.' % file_pattern)

    return match_result.metadata_list


class MatchFiles(beam.PTransform):
  """Matches a file pattern using ``FileSystems.match``.

  This ``PTransform`` returns a ``PCollection`` of matching files in the form
  of ``FileMetadata`` objects."""
  def __init__(
      self,
      file_pattern: str,
      empty_match_treatment=EmptyMatchTreatment.ALLOW_IF_WILDCARD):
    self._file_pattern = file_pattern
    self._empty_match_treatment = empty_match_treatment

  def expand(self, pcoll) -> beam.PCollection[filesystem.FileMetadata]:
    return pcoll.pipeline | beam.Create([self._file_pattern]) | MatchAll(
        empty_match_treatment=self._empty_match_treatment)


class MatchAll(beam.PTransform):
  """Matches file patterns from the input PCollection via ``FileSystems.match``.

  This ``PTransform`` returns a ``PCollection`` of matching files in the form
  of ``FileMetadata`` objects."""
  def __init__(self, empty_match_treatment=EmptyMatchTreatment.ALLOW):
    self._empty_match_treatment = empty_match_treatment

  def expand(
      self,
      pcoll: beam.PCollection,
  ) -> beam.PCollection[filesystem.FileMetadata]:
    return pcoll | beam.ParDo(_MatchAllFn(self._empty_match_treatment))


class ReadableFile(object):
  """A utility class for accessing files."""
  def __init__(self, metadata, compression=None):
    self.metadata = metadata
    self._compression = compression

  def open(self, mime_type='text/plain', compression_type=None):
    compression = (
        compression_type or self._compression or
        filesystems.CompressionTypes.AUTO)
    return filesystems.FileSystems.open(
        self.metadata.path, mime_type=mime_type, compression_type=compression)

  def read(self, mime_type='application/octet-stream'):
    return self.open(mime_type).read()

  def read_utf8(self):
    return self.open().read().decode('utf-8')


class _ReadMatchesFn(beam.DoFn):
  def __init__(self, compression, skip_directories):
    self._compression = compression
    self._skip_directories = skip_directories

  def process(
      self,
      file_metadata: Union[str, filesystem.FileMetadata],
  ) -> Iterable[ReadableFile]:
    metadata = (
        filesystem.FileMetadata(file_metadata, 0) if isinstance(
            file_metadata, str) else file_metadata)

    if ((metadata.path.endswith('/') or metadata.path.endswith('\\')) and
        self._skip_directories):
      return
    elif metadata.path.endswith('/') or metadata.path.endswith('\\'):
      raise BeamIOError(
          'Directories are not allowed in ReadMatches transform.'
          'Found %s.' % metadata.path)

    # TODO: Mime type? Other arguments? Maybe arguments passed in to transform?
    yield ReadableFile(metadata, self._compression)


class _PollClock(object):
  """Shares one clock reading per poll round, so the start gate and the poll
  budget judge the ``start_timestamp`` boundary consistently."""
  def __init__(self):
    self.last_poll_micros: Optional[int] = None


class _WatchWindowTermination(TerminationCondition):
  """Stops after the polls that fall in the ``[start, stop)`` window.

  ``max_polls`` is the ``PeriodicImpulse`` tick count
  ``ceil((stop - start) / interval)``; polls before ``start`` are waiting
  rounds and do not consume the budget.
  """
  def __init__(self, clock: _PollClock, start_micros: int, max_polls: int):
    self._clock = clock
    self._start_micros = start_micros
    self._max_polls = max_polls

  def for_new_input(self, now, element):
    return 0

  def on_poll_complete(self, state):
    poll_micros = self._clock.last_poll_micros
    if poll_micros is not None and poll_micros >= self._start_micros:
      return state + 1
    return state

  def can_stop_polling(self, now, state):
    return state >= self._max_polls

  def state_coder(self):
    return VarIntCoder()


def _ensure_mtime(metadata: filesystem.FileMetadata) -> float:
  # A missing (zero) timestamp is rejected because every file would then carry
  # the same one, and updates could never be told apart.
  if not metadata.last_updated_in_seconds:
    raise BeamIOError(
        'MatchContinuously deduplicates by last-modified time, but %s reports '
        'none.' % metadata.path)
  return metadata.last_updated_in_seconds


def _file_path_key(metadata: filesystem.FileMetadata) -> str:
  return metadata.path


def _file_path_and_mtime_key(
    metadata: filesystem.FileMetadata) -> tuple[str, float]:
  # Keying on the last-modified time makes a changed file look new again.
  return metadata.path, _ensure_mtime(metadata)


class _MatchContinuouslyPollFn(PollFn):
  """Polls a file pattern, honoring empty-match rules.

  A poll before ``start_timestamp`` emits nothing. Matches carry the poll time
  as their event time, or their last-modified time under ``mtime_timestamps``,
  where the watermark trails the newest last-modified time for as long as
  polls keep turning up newer ones.
  """
  def __init__(
      self,
      empty_match_treatment,
      start_timestamp,
      clock=None,
      mtime_timestamps=False):
    self._empty_match_treatment = empty_match_treatment
    self._start_micros = Timestamp.of(start_timestamp).micros
    self._clock = clock if clock is not None else _PollClock()
    self._mtime_timestamps = mtime_timestamps
    # Greatest last-modified time handed out so far, to tell a poll that found
    # something newer from one that only re-listed what was already there.
    self._newest_mtime = None  # type: Optional[Timestamp]

  def __call__(self, file_pattern: str) -> PollResult[filesystem.FileMetadata]:
    now = Timestamp.now()
    self._clock.last_poll_micros = now.micros
    if now.micros < self._start_micros:
      return PollResult.incomplete(())
    match_result = filesystems.FileSystems.match([file_pattern])[0]
    if (not match_result.metadata_list and
        not EmptyMatchTreatment.allow_empty_match(file_pattern,
                                                  self._empty_match_treatment)):
      raise BeamIOError(
          'Empty match for pattern %s. Disallowed.' % file_pattern)
    if not self._mtime_timestamps:
      return PollResult.incomplete(
          match_result.metadata_list, timestamp=now).with_watermark(now)
    outputs = [
        TimestampedValue(metadata, Timestamp.of(_ensure_mtime(metadata)))
        for metadata in match_result.metadata_list
    ]
    # A poll that turned up a newer last-modified time has just read the
    # filesystem clock, so the watermark stops there, capped at the poll time.
    # A poll that found nothing newer takes the poll time, so a quiet
    # directory does not stall event-time windows.
    newest = max((output.timestamp for output in outputs), default=None)
    if newest is not None and (self._newest_mtime is None or
                               newest > self._newest_mtime):
      self._newest_mtime = newest
      watermark = min(newest, now)
    else:
      watermark = now
    return PollResult.incomplete(outputs).with_watermark(watermark)


class MatchContinuously(beam.PTransform):
  """Checks for new files for a given pattern every interval.

  This ``PTransform`` returns a ``PCollection`` of matching files in the form
  of ``FileMetadata`` objects.

  MatchContinuously is experimental.  No backwards-compatibility
  guarantees.

  Deduplication state is checkpointed, so a runner with checkpointing enabled
  restores it after a restart and does not reprocess files. That state grows
  with the number of files matched, unless ``timestamp_cursor`` bounds it. For
  a growing directory on GCS, consider an alternate technique such as Pub/Sub
  Notifications (https://cloud.google.com/storage/docs/pubsub-notifications).

  A match carries the poll time as its event time, and the watermark follows
  the poll time. Under ``timestamp_cursor`` a match carries its last-modified
  time instead, and the watermark holds at the newest one matched, capped at
  the poll time, until a poll turns up nothing newer and releases it.
  """
  def __init__(
      self,
      file_pattern,
      interval=360.0,
      has_deduplication=True,
      start_timestamp=Timestamp.now(),
      stop_timestamp=MAX_TIMESTAMP,
      match_updated_files=False,
      apply_windowing=False,
      empty_match_treatment=EmptyMatchTreatment.ALLOW,
      timestamp_cursor=False):
    """Initializes a MatchContinuously transform.

    Args:
      file_pattern: The file path to read from.
      interval: Interval at which to check for files in seconds.
      has_deduplication: Whether files already read are discarded or not.
      start_timestamp: Timestamp for start file checking.
      stop_timestamp: Timestamp after which no more files will be checked.
      match_updated_files: (When has_deduplication is set to True) whether match
        file with timestamp changes.
      apply_windowing: Whether each element should be assigned to
        individual window. If false, all elements will reside in global window.
      timestamp_cursor: (When match_updated_files and has_deduplication are set
        to True) bound the deduplication state by last-modified time. By
        default, all file modification history is tracked. If set to true, file
        modification history prior to the max(mtime of last poll result) are
        dropped, for better performance. A file that appears with an older
        last-modified time is then taken as already seen and skipped.
    """

    self.file_pattern = file_pattern
    self.interval = interval
    self.has_deduplication = has_deduplication
    self.start_ts = start_timestamp
    self.stop_ts = stop_timestamp
    self.match_upd = match_updated_files
    self.apply_windowing = apply_windowing
    self.empty_match_treatment = empty_match_treatment
    self.timestamp_cursor = timestamp_cursor
    if timestamp_cursor:
      if not has_deduplication:
        raise ValueError(
            'MatchContinuously(timestamp_cursor=True) deduplicates, so it '
            'requires has_deduplication=True.')
      if not match_updated_files:
        _LOGGER.warning(
            'MatchContinuously(timestamp_cursor=True) implies '
            'match_updated_files=True.')
        self.match_upd = True
    else:
      _LOGGER.warning(
          'Matching Continuously is stateful, and can scale poorly. '
          'Consider using Pub/Sub Notifications '
          '(https://cloud.google.com/storage/docs/pubsub-notifications) '
          'if possible')

  def expand(self, pbegin) -> beam.PCollection[filesystem.FileMetadata]:
    if Duration.of(self.interval).micros <= 0:
      raise ValueError('MatchContinuously interval must be positive.')
    if self.has_deduplication:
      match_files = self._match_deduplicated(pbegin)
    else:
      match_files = self._match_all_each_poll(pbegin)

    # Apply windowing last because dedup relies on the global window.
    if self.apply_windowing:
      match_files = match_files | beam.WindowInto(FixedWindows(self.interval))

    return match_files

  def _match_deduplicated(self,
                          pbegin) -> beam.PCollection[filesystem.FileMetadata]:
    # Watch emits each file once per dedup key: the path, joined by the mtime
    # when matching updated files. stop_timestamp bounds the polls to
    # [start, stop).
    clock = _PollClock()
    if self.stop_ts == MAX_TIMESTAMP:
      termination = never()
    else:
      start_ts = Timestamp.of(self.start_ts)
      stop_ts = Timestamp.of(self.stop_ts)
      if stop_ts < start_ts:
        raise ValueError(
            'MatchContinuously stop_timestamp %s precedes start_timestamp %s' %
            (stop_ts, start_ts))
      interval_micros = Duration.of(self.interval).micros
      span_micros = (stop_ts - start_ts).micros
      # Ceiling division reproduces PeriodicImpulse's tick count; the window
      # upper bound is exclusive.
      max_polls = -(-span_micros // interval_micros)
      if max_polls == 0:
        # An empty [start, stop) window never ticks; the impulse path keeps
        # the output empty without Watch's unconditional first poll.
        return self._match_all_each_poll(pbegin)
      termination = _WatchWindowTermination(clock, start_ts.micros, max_polls)
    poll_fn = _MatchContinuouslyPollFn(
        self.empty_match_treatment,
        self.start_ts,
        clock,
        mtime_timestamps=self.timestamp_cursor)
    # The key coder is inferred from the key function's return annotation.
    watch = Watch(
        poll_fn,
        poll_interval=self.interval,
        termination=termination,
        output_key_fn=(
            _file_path_and_mtime_key if self.match_upd else _file_path_key),
        timestamp_cursor=self.timestamp_cursor)
    # Watch emits (pattern, file) pairs; keep the FileMetadata output type so
    # downstream transforms stay typed instead of falling back to Any.
    return (
        pbegin
        | 'Impulse' >> beam.Create([self.file_pattern])
        | 'Watch' >> watch
        | 'DropPattern' >> beam.Map(lambda kv: kv[1]).with_output_types(
            filesystem.FileMetadata))

  def _match_all_each_poll(self,
                           pbegin) -> beam.PCollection[filesystem.FileMetadata]:
    # No deduplication: re-emit every match on each poll.
    return (
        pbegin
        | PeriodicImpulse(
            start_timestamp=self.start_ts,
            stop_timestamp=self.stop_ts,
            fire_interval=self.interval)
        | 'GetFilePattern' >> beam.Map(lambda x: self.file_pattern)
        | MatchAll(self.empty_match_treatment))


class ReadMatches(beam.PTransform):
  """Converts each result of MatchFiles() or MatchAll() to a ReadableFile.

   This helps read in a file's contents or obtain a file descriptor."""
  def __init__(self, compression=None, skip_directories=True):
    self._compression = compression
    self._skip_directories = skip_directories

  def expand(
      self,
      pcoll: beam.PCollection[Union[str, filesystem.FileMetadata]],
  ) -> beam.PCollection[ReadableFile]:
    return pcoll | beam.ParDo(
        _ReadMatchesFn(self._compression, self._skip_directories))


class FileSink(object):
  """Specifies how to write elements to individual files in ``WriteToFiles``.

  A Sink class must implement the following:

   - The ``open`` method, which initializes writing to a file handler (it is not
     responsible for opening the file handler itself).
   - The ``write`` method, which writes an element to the file that was passed
     in ``open``.
   - The ``flush`` method, which flushes any buffered state. This is most often
     called before closing a file (but not exclusively called in that
     situation). The sink is not responsible for closing the file handler.
  A Sink class can override the following:
   - The ``create_metadata`` method, which creates all metadata passed to
     Filesystems.create.
   """
  def create_metadata(
      self, destination: str, full_file_name: str) -> FileMetadata:
    return FileMetadata(
        mime_type="application/octet-stream",
        compression_type=CompressionTypes.AUTO)

  def open(self, fh: BinaryIO) -> None:
    raise NotImplementedError

  def write(self, record):
    raise NotImplementedError

  def flush(self):
    raise NotImplementedError


@beam.typehints.with_input_types(str)
class TextSink(FileSink):
  """A sink that encodes utf8 elements, and writes to file handlers.

  This sink simply calls file_handler.write(record.encode('utf8') + '\n') on all
  records that come into it.
  """
  def open(self, fh):
    self._fh = fh

  def write(self, record):
    self._fh.write(record.encode('utf8'))
    self._fh.write(b'\n')

  def flush(self):
    self._fh.flush()


def prefix_naming(prefix):
  return default_file_naming(prefix)


_DEFAULT_FILE_NAME_TEMPLATE = (
    '{prefix}-{start}-{end}-{pane}-'
    '{shard:05d}-of-{total_shards:05d}'
    '{suffix}{compression}')


def _format_shard(
    window, pane, shard_index, total_shards, compression, prefix, suffix):
  kwargs = {
      'prefix': prefix,
      'start': '',
      'end': '',
      'pane': '',
      'shard': 0,
      'total_shards': 0,
      'suffix': '',
      'compression': ''
  }

  if total_shards is not None and shard_index is not None:
    kwargs['shard'] = int(shard_index)
    kwargs['total_shards'] = int(total_shards)

  if window != GlobalWindow():
    kwargs['start'] = window.start.to_utc_datetime().isoformat()
    kwargs['end'] = window.end.to_utc_datetime().isoformat()

  # TODO(https://github.com/apache/beam/issues/18721): Add support for PaneInfo
  # If the PANE is the ONLY firing in the window, we don't add it.
  #if pane and not (pane.is_first and pane.is_last):
  #  kwargs['pane'] = pane.index

  if suffix:
    kwargs['suffix'] = suffix

  if compression:
    kwargs['compression'] = '.%s' % compression

  # Remove separators for unused template parts.
  format = _DEFAULT_FILE_NAME_TEMPLATE
  if shard_index is None:
    format = format.replace('-{shard:05d}', '')
  if total_shards is None:
    format = format.replace('-of-{total_shards:05d}', '')
  for name, value in kwargs.items():
    if value in (None, ''):
      format = format.replace('-{%s}' % name, '')

  return format.format(**kwargs)


FileNaming = Callable[[Any, Any, int, int, Any, str, str], str]


def destination_prefix_naming(suffix=None) -> FileNaming:
  def _inner(window, pane, shard_index, total_shards, compression, destination):
    prefix = str(destination)
    return _format_shard(
        window, pane, shard_index, total_shards, compression, prefix, suffix)

  return _inner


def default_file_naming(prefix, suffix=None) -> FileNaming:
  def _inner(window, pane, shard_index, total_shards, compression, destination):
    return _format_shard(
        window, pane, shard_index, total_shards, compression, prefix, suffix)

  return _inner


def single_file_naming(prefix, suffix=None) -> FileNaming:
  def _inner(window, pane, shard_index, total_shards, compression, destination):
    assert shard_index in (0, None), shard_index
    assert total_shards in (1, None), total_shards
    return _format_shard(window, pane, None, None, compression, prefix, suffix)

  return _inner


_FileResult = collections.namedtuple(
    'FileResult', [
        'file_name',
        'shard_index',
        'total_shards',
        'window',
        'pane',
        'destination'
    ])


# Adding a class to contain PyDoc.
class FileResult(_FileResult):
  """A descriptor of a file that has been written."""
  pass


class WriteToFiles(beam.PTransform):
  r"""Write the incoming PCollection to a set of output files.

  The incoming ``PCollection`` may be bounded or unbounded.

  **Note:** For unbounded ``PCollection``\s, this transform does not support
  multiple firings per Window (due to the fact that files are named only by
  their destination, and window, at the moment).
  """

  # We allow up to 20 different destinations to be written in a single bundle.
  # Too many files will add memory pressure to the worker, so we let it be 20.
  MAX_NUM_WRITERS_PER_BUNDLE = 20

  DEFAULT_SHARDING = 5

  def __init__(
      self,
      path,
      file_naming=None,
      destination=None,
      temp_directory=None,
      sink=None,
      shards=None,
      output_fn=None,
      max_writers_per_bundle=MAX_NUM_WRITERS_PER_BUNDLE):
    """Initializes a WriteToFiles transform.

    Args:
      path (str, ValueProvider): The directory to write files into.
      file_naming (callable): A callable that takes in a window, pane,
        shard_index, total_shards and compression; and returns a file name.
      destination (callable): If this argument is provided, the sink parameter
        must also be a callable.
      temp_directory (str, ValueProvider): To ensure atomicity in the transform,
        the output is written into temporary files, which are written to a
        directory that is meant to be temporary as well. Once the whole output
        has been written, the files are moved into their final destination, and
        given their final names. By default, the temporary directory will be
        within the temp_location of your pipeline.
      sink (callable, ~apache_beam.io.fileio.FileSink): The sink to use to write
        into a file. It should implement the methods of a ``FileSink``. Pass a
        class signature or an instance of FileSink to this parameter. If none is
        provided, a ``TextSink`` is used.
      shards (int): The number of shards per destination and trigger firing.
      output_fn (callable, optional): A callable to process the output. This
        parameter is currently unused and retained for backward compatibility.
      max_writers_per_bundle (int): The number of writers that can be open
        concurrently in a single worker that's processing one bundle.
    """
    self.path = (
        path if isinstance(path, ValueProvider) else StaticValueProvider(
            str, path))
    self.file_naming_fn = file_naming or default_file_naming('output')
    self.destination_fn = self._get_destination_fn(destination)
    self._temp_directory = temp_directory
    self.sink_fn = self._get_sink_fn(sink)
    self.shards = shards or WriteToFiles.DEFAULT_SHARDING
    self.output_fn = output_fn or (lambda x: x)

    self._max_num_writers_per_bundle = max_writers_per_bundle

  @staticmethod
  def _get_sink_fn(input_sink) -> Callable[[Any], FileSink]:
    if isinstance(input_sink, type) and issubclass(input_sink, FileSink):
      return lambda x: input_sink()
    elif isinstance(input_sink, FileSink):
      kls = input_sink.__class__
      return lambda x: kls()
    elif callable(input_sink):
      return input_sink
    else:
      return lambda x: TextSink()

  @staticmethod
  def _get_destination_fn(destination) -> Callable[[Any], str]:
    if isinstance(destination, ValueProvider):
      return lambda elm: destination.get()
    elif callable(destination):
      return destination
    else:
      return lambda elm: destination

  def expand(self, pcoll):
    p = pcoll.pipeline

    if not self._temp_directory:
      temp_location = (
          p.options.view_as(GoogleCloudOptions).temp_location or
          self.path.get())
      dir_uid = str(uuid.uuid4())
      self._temp_directory = StaticValueProvider(
          str, filesystems.FileSystems.join(temp_location, '.temp%s' % dir_uid))
      _LOGGER.info('Added temporary directory %s', self._temp_directory.get())

    output = (
        pcoll
        | beam.ParDo(
            _WriteUnshardedRecordsFn(
                base_path=self._temp_directory,
                destination_fn=self.destination_fn,
                sink_fn=self.sink_fn,
                max_writers_per_bundle=self._max_num_writers_per_bundle)).
        with_outputs(
            _WriteUnshardedRecordsFn.SPILLED_RECORDS,
            _WriteUnshardedRecordsFn.WRITTEN_FILES))

    written_files_pc = output[_WriteUnshardedRecordsFn.WRITTEN_FILES]
    spilled_records_pc = output[_WriteUnshardedRecordsFn.SPILLED_RECORDS]

    more_written_files_pc = (
        spilled_records_pc
        | beam.ParDo(
            _AppendShardedDestination(self.destination_fn, self.shards))
        | "GroupRecordsByDestinationAndShard" >> beam.GroupByKey()
        | beam.ParDo(
            _WriteShardedRecordsFn(
                self._temp_directory, self.sink_fn, self.shards)))

    files_by_destination_pc = (
        (written_files_pc, more_written_files_pc)
        | beam.Flatten()
        | beam.Map(lambda file_result: (file_result.destination, file_result))
        | "GroupTempFilesByDestination" >> beam.GroupByKey())

    # Now we should take the temporary files, and write them to the final
    # destination, with their proper names.

    file_results = (
        files_by_destination_pc
        | beam.ParDo(
            _MoveTempFilesIntoFinalDestinationFn(
                self.path, self.file_naming_fn, self._temp_directory)))

    return file_results


def _create_writer(
    base_path,
    writer_key: tuple[str, IntervalWindow],
    create_metadata_fn: CreateFileMetadataFn,
):
  try:
    filesystems.FileSystems.mkdirs(base_path)
  except IOError:
    # Directory already exists.
    pass

  destination = writer_key[0]

  # The file name has a prefix determined by destination+window, along with
  # a random string. This allows us to retrieve orphaned files later on.
  file_name = '%s_%s' % (abs(hash(writer_key)), uuid.uuid4())
  full_file_name = filesystems.FileSystems.join(base_path, file_name)
  metadata = create_metadata_fn(destination, full_file_name)
  return full_file_name, filesystems.FileSystems.create(
      full_file_name,
      **metadata._asdict())


class _MoveTempFilesIntoFinalDestinationFn(beam.DoFn):
  def __init__(self, path, file_naming_fn, temp_dir):
    self.path = path
    self.file_naming_fn = file_naming_fn
    self.temporary_directory = temp_dir

  def process(self, element, w=beam.DoFn.WindowParam):
    destination = element[0]
    # list of FileResult objects for temp files
    temp_file_results = list(element[1])
    # list of FileResult objects for final files
    final_file_results = []

    for i, r in enumerate(temp_file_results):
      # TODO(pabloem): Handle compression for files.
      final_file_name = self.file_naming_fn(
          r.window, r.pane, i, len(temp_file_results), '', destination)

      final_file_results.append(
          FileResult(
              final_file_name,
              i,
              len(temp_file_results),
              r.window,
              r.pane,
              destination))

    move_from = [f.file_name for f in temp_file_results]
    move_to = [f.file_name for f in final_file_results]

    _LOGGER.info(
        'Moving %d temporary files to dir: %s as %s',
        len(move_from),
        self.path.get(),
        move_to)

    try:
      filesystems.FileSystems.mkdirs(self.path.get())
    except IOError as e:
      cause = repr(e)
      if 'FileExistsError' not in cause:
        # Usually harmless. Especially if see FileExistsError so no need to log
        _LOGGER.debug('Fail to create dir for final destination: %s', cause)

    try:
      filesystems.FileSystems.rename(
          move_from,
          [filesystems.FileSystems.join(self.path.get(), f) for f in move_to])
    except BeamIOError:
      # This error is not serious, because it may happen on a retry of the
      # bundle. We simply log it.
      _LOGGER.debug(
          'Exception occurred during moving files: %s. This may be due to a'
          ' bundle being retried.',
          move_from)

    yield from final_file_results

    _LOGGER.debug(
        'Checking orphaned temporary files for destination %s and window %s',
        destination,
        w)
    writer_key = (destination, w)
    self._check_orphaned_files(writer_key)

  def _check_orphaned_files(self, writer_key):
    try:
      prefix = filesystems.FileSystems.join(
          self.temporary_directory.get(), str(abs(hash(writer_key))))
      match_result = filesystems.FileSystems.match(['%s*' % prefix])
      orphaned_files = [m.path for m in match_result[0].metadata_list]

      if len(orphaned_files) > 0:
        _LOGGER.warning(
            'Some files may be left orphaned in the temporary folder: %s. '
            'This may be a result of retried work items or insufficient'
            'permissions to delete these temp files.',
            orphaned_files)
    except BeamIOError as e:
      _LOGGER.warning('Exceptions when checking orphaned files: %s', e)


class _WriteShardedRecordsFn(beam.DoFn):
  def __init__(
      self, base_path, sink_fn: Callable[[Any], FileSink], shards: int):
    self.base_path = base_path
    self.sink_fn = sink_fn
    self.shards = shards

  def process(
      self, element, w=beam.DoFn.WindowParam, pane=beam.DoFn.PaneInfoParam):
    destination_and_shard = element[0]
    destination = destination_and_shard[0]
    shard = destination_and_shard[1]
    records = element[1]

    sink = self.sink_fn(destination)

    full_file_name, writer = _create_writer(
        base_path=self.base_path.get(),
        writer_key=(destination, w),
        create_metadata_fn=sink.create_metadata)

    sink.open(writer)

    for r in records:
      sink.write(r)

    sink.flush()
    writer.close()

    _LOGGER.info(
        'Writing file %s for destination %s and shard %s',
        full_file_name,
        destination,
        repr(shard))

    yield FileResult(
        full_file_name,
        shard_index=shard,
        total_shards=self.shards,
        window=w,
        pane=pane,
        destination=destination)


class _AppendShardedDestination(beam.DoFn):
  def __init__(self, destination: Callable[[Any], str], shards: int):
    self.destination_fn = destination
    self.shards = shards

    # We start the shards for a single destination at an arbitrary point.
    self._shard_counter: collections.defaultdict[
        str,
        int] = collections.defaultdict(lambda: random.randrange(self.shards))

  def _next_shard_for_destination(self, destination):
    self._shard_counter[destination] = ((self._shard_counter[destination] + 1) %
                                        self.shards)

    return self._shard_counter[destination]

  def process(self, record):
    destination = self.destination_fn(record)
    shard = self._next_shard_for_destination(destination)

    yield ((destination, shard), record)


class _WriteUnshardedRecordsFn(beam.DoFn):

  SPILLED_RECORDS = 'spilled_records'
  WRITTEN_FILES = 'written_files'

  _writers_and_sinks: dict[tuple[str, BoundedWindow], tuple[BinaryIO,
                                                            FileSink]] = None
  _file_names: dict[tuple[str, BoundedWindow], str] = None

  def __init__(
      self,
      base_path,
      destination_fn,
      sink_fn,
      max_writers_per_bundle=WriteToFiles.MAX_NUM_WRITERS_PER_BUNDLE):
    self.base_path = base_path
    self.destination_fn = destination_fn
    self.sink_fn = sink_fn
    self.max_num_writers_per_bundle = max_writers_per_bundle

  def start_bundle(self):
    self._writers_and_sinks = {}
    self._file_names = {}

  def process(
      self, record, w=beam.DoFn.WindowParam, pane=beam.DoFn.PaneInfoParam):
    destination = self.destination_fn(record)

    writer, sink = self._get_or_create_writer_and_sink(destination, w)

    if not writer:
      return [beam.pvalue.TaggedOutput(self.SPILLED_RECORDS, record)]
    else:
      sink.write(record)

  def _get_or_create_writer_and_sink(self, destination, window):
    """Returns a tuple of writer, sink."""
    writer_key = (destination, window)
    if writer_key in self._writers_and_sinks:
      return self._writers_and_sinks.get(writer_key)
    elif len(self._writers_and_sinks) >= self.max_num_writers_per_bundle:
      # The writer does not exist, and we have too many writers already.
      return None, None
    else:
      # The writer does not exist, but we can still create a new one.
      sink = self.sink_fn(destination)

      full_file_name, writer = _create_writer(
          base_path=self.base_path.get(),
          writer_key=writer_key,
          create_metadata_fn=sink.create_metadata)

      sink.open(writer)
      self._writers_and_sinks[writer_key] = (writer, sink)
      self._file_names[writer_key] = full_file_name
      return self._writers_and_sinks[writer_key]

  def finish_bundle(self):
    for key, (writer, sink) in self._writers_and_sinks.items():

      sink.flush()
      writer.close()

      file_result = FileResult(
          self._file_names[key],
          shard_index=-1,
          total_shards=0,
          window=key[1],
          pane=None,  # TODO(pabloem): get the pane info
          destination=key[0])

      yield beam.pvalue.TaggedOutput(
          self.WRITTEN_FILES,
          beam.transforms.window.WindowedValue(
              file_result,
              timestamp=key[1].start,
              windows=[key[1]]  # TODO(pabloem) HOW DO WE GET THE PANE
          ))