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

"""Windowing concepts.

A WindowInto transform logically divides up or groups the elements of a
PCollection into finite windows according to a windowing function (derived from
WindowFn).

The output of WindowInto contains the same elements as input, but they have been
logically assigned to windows. The next GroupByKey(s) transforms, including one
within a composite transform, will group by the combination of keys and windows.

Windowing a PCollection allows chunks of it to be processed individually, before
the entire PCollection is available.  This is especially important for
PCollection(s) with unbounded size, since the full PCollection is never
available at once, since more data is continually arriving. For PCollection(s)
with a bounded size (aka. conventional batch mode), by default, all data is
implicitly in a single window (see GlobalWindows), unless WindowInto is
applied.

For example, a simple form of windowing divides up the data into fixed-width
time intervals, using FixedWindows.

Seconds are used as the time unit for the built-in windowing primitives here.
Integer or floating point seconds can be passed to these primitives.

Internally, seconds, with microsecond granularity, are stored as
timeutil.Timestamp and timeutil.Duration objects. This is done to avoid
precision errors that would occur with floating point representations.

Custom windowing function classes can be created, by subclassing from
WindowFn.
"""

# pytype: skip-file

import abc
from collections.abc import Iterable
from functools import total_ordering
from typing import Any
from typing import Generic
from typing import Optional
from typing import TypeVar

from google.protobuf import duration_pb2
from google.protobuf import timestamp_pb2

from apache_beam.coders import coders
from apache_beam.portability import common_urns
from apache_beam.portability import python_urns
from apache_beam.portability.api import beam_runner_api_pb2
from apache_beam.portability.api import standard_window_fns_pb2
from apache_beam.transforms import timeutil
from apache_beam.utils import proto_utils
from apache_beam.utils import urns
from apache_beam.utils import windowed_value
from apache_beam.utils.timestamp import MIN_TIMESTAMP
from apache_beam.utils.timestamp import Duration
from apache_beam.utils.timestamp import DurationTypes  # pylint: disable=unused-import
from apache_beam.utils.timestamp import Timestamp
from apache_beam.utils.timestamp import TimestampTypes  # pylint: disable=unused-import
from apache_beam.utils.windowed_value import WindowedValue

__all__ = [
    'TimestampCombiner',
    'WindowFn',
    'BoundedWindow',
    'IntervalWindow',
    'TimestampedValue',
    'GlobalWindow',
    'NonMergingWindowFn',
    'GlobalWindows',
    'FixedWindows',
    'SlidingWindows',
    'Sessions',
]


# TODO(ccy): revisit naming and semantics once Java Apache Beam finalizes their
# behavior.
class TimestampCombiner(object):
  """Determines how output timestamps of grouping operations are assigned."""

  OUTPUT_AT_EOW = beam_runner_api_pb2.OutputTime.END_OF_WINDOW
  OUTPUT_AT_EARLIEST = beam_runner_api_pb2.OutputTime.EARLIEST_IN_PANE
  OUTPUT_AT_LATEST = beam_runner_api_pb2.OutputTime.LATEST_IN_PANE
  # TODO(robertwb): Add this to the runner API or remove it.
  OUTPUT_AT_EARLIEST_TRANSFORMED = 'OUTPUT_AT_EARLIEST_TRANSFORMED'

  @staticmethod
  def get_impl(
      timestamp_combiner: beam_runner_api_pb2.OutputTime.Enum,
      window_fn: 'WindowFn') -> timeutil.TimestampCombinerImpl:
    if timestamp_combiner == TimestampCombiner.OUTPUT_AT_EOW:
      return timeutil.OutputAtEndOfWindowImpl()
    elif timestamp_combiner == TimestampCombiner.OUTPUT_AT_EARLIEST:
      return timeutil.OutputAtEarliestInputTimestampImpl()
    elif timestamp_combiner == TimestampCombiner.OUTPUT_AT_LATEST:
      return timeutil.OutputAtLatestInputTimestampImpl()
    elif timestamp_combiner == TimestampCombiner.OUTPUT_AT_EARLIEST_TRANSFORMED:
      return timeutil.OutputAtEarliestTransformedInputTimestampImpl(window_fn)
    else:
      raise ValueError('Invalid TimestampCombiner: %s.' % timestamp_combiner)


class WindowFn(urns.RunnerApiFn, metaclass=abc.ABCMeta):
  """An abstract windowing function defining a basic assign and merge."""
  class AssignContext(object):
    """Context passed to WindowFn.assign()."""
    def __init__(
        self,
        timestamp: TimestampTypes,
        element: Optional[Any] = None,
        window: Optional['BoundedWindow'] = None) -> None:
      self.timestamp = Timestamp.of(timestamp)
      self.element = element
      self.window = window

  @abc.abstractmethod
  def assign(self,
             assign_context: 'AssignContext') -> Iterable['BoundedWindow']:
    # noqa: F821

    """Associates windows to an element.

    Arguments:
      assign_context: Instance of AssignContext.

    Returns:
      An iterable of BoundedWindow.
    """
    raise NotImplementedError

  class MergeContext(object):
    """Context passed to WindowFn.merge() to perform merging, if any."""
    def __init__(self, windows: Iterable['BoundedWindow']) -> None:
      self.windows = list(windows)

    def merge(
        self,
        to_be_merged: Iterable['BoundedWindow'],
        merge_result: 'BoundedWindow') -> None:
      raise NotImplementedError

  @abc.abstractmethod
  def merge(self, merge_context: 'WindowFn.MergeContext') -> None:
    """Returns a window that is the result of merging a set of windows."""
    raise NotImplementedError

  def is_merging(self) -> bool:
    """Returns whether this WindowFn merges windows."""
    return True

  @abc.abstractmethod
  def get_window_coder(self) -> coders.Coder:
    raise NotImplementedError

  def get_transformed_output_time(
      self, window: 'BoundedWindow', input_timestamp: Timestamp) -> Timestamp:  # pylint: disable=unused-argument
    """Given input time and output window, returns output time for window.

    If TimestampCombiner.OUTPUT_AT_EARLIEST_TRANSFORMED is used in the
    Windowing, the output timestamp for the given window will be the earliest
    of the timestamps returned by get_transformed_output_time() for elements
    of the window.

    Arguments:
      window: Output window of element.
      input_timestamp: Input timestamp of element as a timeutil.Timestamp
        object.

    Returns:
      Transformed timestamp.
    """
    # By default, just return the input timestamp.
    return input_timestamp

  urns.RunnerApiFn.register_pickle_urn(python_urns.PICKLED_WINDOWFN)


class BoundedWindow(object):
  """A window for timestamps in range (-infinity, end).

  Attributes:
    end: End of window.
  """
  def __init__(self, end: TimestampTypes) -> None:
    self._end = Timestamp.of(end)

  @property
  def start(self) -> Timestamp:
    raise NotImplementedError

  @property
  def end(self) -> Timestamp:
    return self._end

  def max_timestamp(self) -> Timestamp:
    return self.end.predecessor()

  def __eq__(self, other):
    raise NotImplementedError

  def __ne__(self, other):
    #  Order first by endpoint, then arbitrarily
    return self.end != other.end or hash(self) != hash(other)

  def __lt__(self, other):
    if self.end != other.end:
      return self.end < other.end
    return hash(self) < hash(other)

  def __le__(self, other):
    if self.end != other.end:
      return self.end <= other.end
    return hash(self) <= hash(other)

  def __gt__(self, other):
    if self.end != other.end:
      return self.end > other.end
    return hash(self) > hash(other)

  def __ge__(self, other):
    if self.end != other.end:
      return self.end >= other.end
    return hash(self) >= hash(other)

  def __hash__(self):
    raise NotImplementedError

  def __repr__(self):
    return '[?, %s)' % float(self.end)


@total_ordering
class IntervalWindow(windowed_value._IntervalWindowBase, BoundedWindow):
  """A window for timestamps in range [start, end).

  Attributes:
    start: Start of window as seconds since Unix epoch.
    end: End of window as seconds since Unix epoch.
  """
  def __lt__(self, other):
    if self.end != other.end:
      return self.end < other.end
    return hash(self) < hash(other)

  def intersects(self, other: 'IntervalWindow') -> bool:
    return other.start < self.end or self.start < other.end

  def union(self, other: 'IntervalWindow') -> 'IntervalWindow':
    return IntervalWindow(
        min(self.start, other.start), max(self.end, other.end))


V = TypeVar("V")


@total_ordering
class TimestampedValue(Generic[V]):
  """A timestamped value having a value and a timestamp.

  Attributes:
    value: The underlying value.
    timestamp: Timestamp associated with the value as seconds since Unix epoch.
  """
  def __init__(self, value: V, timestamp: TimestampTypes) -> None:
    self.value = value
    self.timestamp = Timestamp.of(timestamp)

  def __eq__(self, other):
    return (
        type(self) == type(other) and self.value == other.value and
        self.timestamp == other.timestamp)

  def __hash__(self):
    return hash((self.value, self.timestamp))

  def __lt__(self, other):
    if type(self) != type(other):
      return type(self).__name__ < type(other).__name__
    if self.value != other.value:
      return self.value < other.value
    return self.timestamp < other.timestamp


class GlobalWindow(BoundedWindow):
  """The default window into which all data is placed (via GlobalWindows)."""
  _instance: Optional['GlobalWindow'] = None

  def __new__(cls):
    if cls._instance is None:
      cls._instance = super(GlobalWindow, cls).__new__(cls)
    return cls._instance

  def __init__(self) -> None:
    super().__init__(GlobalWindow._getTimestampFromProto())

  def __repr__(self):
    return 'GlobalWindow'

  def __hash__(self):
    return hash(type(self))

  def __eq__(self, other):
    # Global windows are always and only equal to each other.
    return self is other or type(self) is type(other)

  @property
  def start(self) -> Timestamp:
    return MIN_TIMESTAMP

  @staticmethod
  def _getTimestampFromProto() -> Timestamp:
    ts_millis = int(
        common_urns.constants.GLOBAL_WINDOW_MAX_TIMESTAMP_MILLIS.constant)
    return Timestamp(micros=ts_millis * 1000)


class NonMergingWindowFn(WindowFn):
  def is_merging(self) -> bool:
    return False

  def merge(self, merge_context: WindowFn.MergeContext) -> None:
    pass  # No merging.


class GlobalWindows(NonMergingWindowFn):
  """A windowing function that assigns everything to one global window."""
  @classmethod
  def windowed_batch(
      cls,
      batch: Any,
      timestamp: Timestamp = MIN_TIMESTAMP,
      pane_info: windowed_value.PaneInfo = windowed_value.PANE_INFO_UNKNOWN
  ) -> windowed_value.WindowedBatch:
    return windowed_value.HomogeneousWindowedBatch.of(
        batch, timestamp, (GlobalWindow(), ), pane_info)

  @classmethod
  def windowed_value(
      cls,
      value: Any,
      timestamp: Timestamp = MIN_TIMESTAMP,
      pane_info: windowed_value.PaneInfo = windowed_value.PANE_INFO_UNKNOWN
  ) -> WindowedValue:
    return WindowedValue(value, timestamp, (GlobalWindow(), ), pane_info)

  @classmethod
  def windowed_value_at_end_of_window(cls, value):
    return cls.windowed_value(value, GlobalWindow().max_timestamp())

  def assign(self,
             assign_context: WindowFn.AssignContext) -> list[GlobalWindow]:
    return [GlobalWindow()]

  def get_window_coder(self) -> coders.GlobalWindowCoder:
    return coders.GlobalWindowCoder()

  def __hash__(self):
    return hash(type(self))

  def __eq__(self, other):
    # Global windowfn is always and only equal to each other.
    return self is other or type(self) is type(other)

  def to_runner_api_parameter(self, context):
    return common_urns.global_windows.urn, None

  @staticmethod
  @urns.RunnerApiFn.register_urn(common_urns.global_windows.urn, None)
  def from_runner_api_parameter(
      unused_fn_parameter, unused_context) -> 'GlobalWindows':
    return GlobalWindows()


class FixedWindows(NonMergingWindowFn):
  """A windowing function that assigns each element to one time interval.

  The attributes size and offset determine in what time interval a timestamp
  will be slotted. The time intervals have the following formula:
  [N * size + offset, (N + 1) * size + offset)

  Attributes:
    size: Size of the window as seconds.
    offset: Offset of this window as seconds. Windows start at
      t=N * size + offset where t=0 is the UNIX epoch. The offset must be a
      value in range [0, size). If it is not it will be normalized to this
      range.
  """
  def __init__(self, size: DurationTypes, offset: TimestampTypes = 0):
    """Initialize a ``FixedWindows`` function for a given size and offset.

    Args:
      size (int): Size of the window in seconds.
      offset(int): Offset of this window as seconds. Windows start at
        t=N * size + offset where t=0 is the UNIX epoch. The offset must be a
        value in range [0, size). If it is not it will be normalized to this
        range.
    """
    if size <= 0:
      raise ValueError('The size parameter must be strictly positive.')
    self.size = Duration.of(size)
    self.offset = Timestamp.of(offset) % self.size

  def assign(self, context: WindowFn.AssignContext) -> list[IntervalWindow]:
    timestamp = context.timestamp
    start = timestamp - (timestamp - self.offset) % self.size
    return [IntervalWindow(start, start + self.size)]

  def get_window_coder(self) -> coders.IntervalWindowCoder:
    return coders.IntervalWindowCoder()

  def __eq__(self, other):
    if type(self) == type(other) == FixedWindows:
      return self.size == other.size and self.offset == other.offset

  def __hash__(self):
    return hash((self.size, self.offset))

  def to_runner_api_parameter(self, context):
    return (
        common_urns.fixed_windows.urn,
        standard_window_fns_pb2.FixedWindowsPayload(
            size=proto_utils.from_micros(
                duration_pb2.Duration, self.size.micros),
            offset=proto_utils.from_micros(
                timestamp_pb2.Timestamp, self.offset.micros)))

  @staticmethod
  @urns.RunnerApiFn.register_urn(
      common_urns.fixed_windows.urn,
      standard_window_fns_pb2.FixedWindowsPayload)
  def from_runner_api_parameter(fn_parameter, unused_context) -> 'FixedWindows':
    return FixedWindows(
        size=Duration(micros=proto_utils.to_micros(fn_parameter.size)),
        offset=Timestamp(micros=proto_utils.to_micros(fn_parameter.offset)))


class SlidingWindows(NonMergingWindowFn):
  """A windowing function that assigns each element to a set of sliding windows.

  The attributes size and offset determine in what time interval a timestamp
  will be slotted. The time intervals have the following formula:
  [N * period + offset, N * period + offset + size)

  Attributes:
    size: Size of the window as seconds.
    period: Period of the windows as seconds.
    offset: Offset of this window as seconds since Unix epoch. Windows start at
      t=N * period + offset where t=0 is the epoch. The offset must be a value
      in range [0, period). If it is not it will be normalized to this range.
  """
  def __init__(
      self,
      size: DurationTypes,
      period: DurationTypes,
      offset: TimestampTypes = 0,
  ):
    if size <= 0:
      raise ValueError('The size parameter must be strictly positive.')
    self.size = Duration.of(size)
    self.period = Duration.of(period)
    self.offset = Timestamp.of(offset) % period

  def assign(self, context: WindowFn.AssignContext) -> list[IntervalWindow]:
    timestamp = context.timestamp
    start = timestamp - ((timestamp - self.offset) % self.period)
    return [
        IntervalWindow(
            (interval_start := Timestamp(micros=s)),
            interval_start + self.size,
        ) for s in range(
            start.micros, timestamp.micros -
            self.size.micros, -self.period.micros)
    ]

  def get_window_coder(self) -> coders.IntervalWindowCoder:
    return coders.IntervalWindowCoder()

  def __eq__(self, other):
    if type(self) == type(other) == SlidingWindows:
      return (
          self.size == other.size and self.offset == other.offset and
          self.period == other.period)

  def __hash__(self):
    return hash((self.offset, self.period))

  def to_runner_api_parameter(self, context):
    return (
        common_urns.sliding_windows.urn,
        standard_window_fns_pb2.SlidingWindowsPayload(
            size=proto_utils.from_micros(
                duration_pb2.Duration, self.size.micros),
            offset=proto_utils.from_micros(
                timestamp_pb2.Timestamp, self.offset.micros),
            period=proto_utils.from_micros(
                duration_pb2.Duration, self.period.micros)))

  @staticmethod
  @urns.RunnerApiFn.register_urn(
      common_urns.sliding_windows.urn,
      standard_window_fns_pb2.SlidingWindowsPayload)
  def from_runner_api_parameter(
      fn_parameter, unused_context) -> 'SlidingWindows':
    return SlidingWindows(
        size=Duration(micros=proto_utils.to_micros(fn_parameter.size)),
        offset=Timestamp(micros=proto_utils.to_micros(fn_parameter.offset)),
        period=Duration(micros=proto_utils.to_micros(fn_parameter.period)))


class Sessions(WindowFn):
  """A windowing function that groups elements into sessions.

  A session is defined as a series of consecutive events
  separated by a specified gap size.

  Attributes:
    gap_size: Size of the gap between windows as floating-point seconds.
  """
  def __init__(self, gap_size: DurationTypes) -> None:
    if gap_size <= 0:
      raise ValueError('The size parameter must be strictly positive.')
    self.gap_size = Duration.of(gap_size)

  def assign(self, context: WindowFn.AssignContext) -> list[IntervalWindow]:
    timestamp = context.timestamp
    return [IntervalWindow(timestamp, timestamp + self.gap_size)]

  def get_window_coder(self) -> coders.IntervalWindowCoder:
    return coders.IntervalWindowCoder()

  def merge(self, merge_context: WindowFn.MergeContext) -> None:
    to_merge: list[BoundedWindow] = []
    end = MIN_TIMESTAMP
    for w in sorted(merge_context.windows, key=lambda w: w.start):
      if to_merge:
        if end > w.start:
          to_merge.append(w)
          if w.end > end:
            end = w.end
        else:
          if len(to_merge) > 1:
            merge_context.merge(
                to_merge, IntervalWindow(to_merge[0].start, end))
          to_merge = [w]
          end = w.end
      else:
        to_merge = [w]
        end = w.end
    if len(to_merge) > 1:
      merge_context.merge(to_merge, IntervalWindow(to_merge[0].start, end))

  def __eq__(self, other):
    if type(self) == type(other) == Sessions:
      return self.gap_size == other.gap_size

  def __hash__(self):
    return hash(self.gap_size)

  def to_runner_api_parameter(self, context):
    return (
        common_urns.session_windows.urn,
        standard_window_fns_pb2.SessionWindowsPayload(
            gap_size=proto_utils.from_micros(
                duration_pb2.Duration, self.gap_size.micros)))

  @staticmethod
  @urns.RunnerApiFn.register_urn(
      common_urns.session_windows.urn,
      standard_window_fns_pb2.SessionWindowsPayload)
  def from_runner_api_parameter(fn_parameter, unused_context) -> 'Sessions':
    return Sessions(
        gap_size=Duration(micros=proto_utils.to_micros(fn_parameter.gap_size)))