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

import operator
import time
from itertools import chain
from datetime import datetime
from typing import (
    Any,
    Callable,
    Generic,
    Hashable,
    Iterable,
    List,
    Optional,
    Tuple,
    TypeVar,
    Union,
    TYPE_CHECKING,
    cast,
    overload,
)

from py4j.protocol import Py4JJavaError
from py4j.java_gateway import JavaObject

from pyspark.storagelevel import StorageLevel
from pyspark.streaming.util import rddToFileName, TransformFunction
from pyspark.core.rdd import portable_hash, RDD
from pyspark.resultiterable import ResultIterable

if TYPE_CHECKING:
    from pyspark.serializers import Serializer
    from pyspark.streaming.context import StreamingContext

__all__ = ["DStream"]

S = TypeVar("S")
T = TypeVar("T")
T_co = TypeVar("T_co", covariant=True)
U = TypeVar("U")
K = TypeVar("K", bound=Hashable)
V = TypeVar("V")


class DStream(Generic[T_co]):
    """
    A Discretized Stream (DStream), the basic abstraction in Spark Streaming,
    is a continuous sequence of RDDs (of the same type) representing a
    continuous stream of data (see :class:`RDD` in the Spark core documentation
    for more details on RDDs).

    DStreams can either be created from live data (such as, data from TCP
    sockets, etc.) using a :class:`StreamingContext` or it can be
    generated by transforming existing DStreams using operations such as
    `map`, `window` and `reduceByKeyAndWindow`. While a Spark Streaming
    program is running, each DStream periodically generates a RDD, either
    from live data or by transforming the RDD generated by a parent DStream.

    DStreams internally is characterized by a few basic properties:
     - A list of other DStreams that the DStream depends on
     - A time interval at which the DStream generates an RDD
     - A function that is used to generate an RDD after each time interval
    """

    def __init__(
        self,
        jdstream: JavaObject,
        ssc: "StreamingContext",
        jrdd_deserializer: "Serializer",
    ):
        self._jdstream = jdstream
        self._ssc = ssc
        self._sc = ssc._sc
        self._jrdd_deserializer = jrdd_deserializer
        self.is_cached = False
        self.is_checkpointed = False

    def context(self) -> "StreamingContext":
        """
        Return the StreamingContext associated with this DStream
        """
        return self._ssc

    def count(self) -> "DStream[int]":
        """
        Return a new DStream in which each RDD has a single element
        generated by counting each RDD of this DStream.
        """
        return self.mapPartitions(lambda i: [sum(1 for _ in i)]).reduce(operator.add)

    def filter(self: "DStream[T]", f: Callable[[T], bool]) -> "DStream[T]":
        """
        Return a new DStream containing only the elements that satisfy predicate.
        """

        def func(iterator: Iterable[T]) -> Iterable[T]:
            return filter(f, iterator)

        return self.mapPartitions(func, True)

    def flatMap(
        self: "DStream[T]",
        f: Callable[[T], Iterable[U]],
        preservesPartitioning: bool = False,
    ) -> "DStream[U]":
        """
        Return a new DStream by applying a function to all elements of
        this DStream, and then flattening the results
        """

        def func(s: int, iterator: Iterable[T]) -> Iterable[U]:
            return chain.from_iterable(map(f, iterator))

        return self.mapPartitionsWithIndex(func, preservesPartitioning)

    def map(
        self: "DStream[T]", f: Callable[[T], U], preservesPartitioning: bool = False
    ) -> "DStream[U]":
        """
        Return a new DStream by applying a function to each element of DStream.
        """

        def func(iterator: Iterable[T]) -> Iterable[U]:
            return map(f, iterator)

        return self.mapPartitions(func, preservesPartitioning)

    def mapPartitions(
        self: "DStream[T]",
        f: Callable[[Iterable[T]], Iterable[U]],
        preservesPartitioning: bool = False,
    ) -> "DStream[U]":
        """
        Return a new DStream in which each RDD is generated by applying
        mapPartitions() to each RDDs of this DStream.
        """

        def func(s: int, iterator: Iterable[T]) -> Iterable[U]:
            return f(iterator)

        return self.mapPartitionsWithIndex(func, preservesPartitioning)

    def mapPartitionsWithIndex(
        self: "DStream[T]",
        f: Callable[[int, Iterable[T]], Iterable[U]],
        preservesPartitioning: bool = False,
    ) -> "DStream[U]":
        """
        Return a new DStream in which each RDD is generated by applying
        mapPartitionsWithIndex() to each RDDs of this DStream.
        """
        return self.transform(lambda rdd: rdd.mapPartitionsWithIndex(f, preservesPartitioning))

    def reduce(self: "DStream[T]", func: Callable[[T, T], T]) -> "DStream[T]":
        """
        Return a new DStream in which each RDD has a single element
        generated by reducing each RDD of this DStream.
        """
        return self.map(lambda x: (None, x)).reduceByKey(func, 1).map(lambda x: x[1])

    def reduceByKey(
        self: "DStream[Tuple[K, V]]",
        func: Callable[[V, V], V],
        numPartitions: Optional[int] = None,
    ) -> "DStream[Tuple[K, V]]":
        """
        Return a new DStream by applying reduceByKey to each RDD.
        """
        if numPartitions is None:
            numPartitions = self._sc.defaultParallelism
        return self.combineByKey(lambda x: x, func, func, numPartitions)

    def combineByKey(
        self: "DStream[Tuple[K, V]]",
        createCombiner: Callable[[V], U],
        mergeValue: Callable[[U, V], U],
        mergeCombiners: Callable[[U, U], U],
        numPartitions: Optional[int] = None,
    ) -> "DStream[Tuple[K, U]]":
        """
        Return a new DStream by applying combineByKey to each RDD.
        """
        if numPartitions is None:
            numPartitions = self._sc.defaultParallelism

        def func(rdd: RDD[Tuple[K, V]]) -> RDD[Tuple[K, U]]:
            return rdd.combineByKey(createCombiner, mergeValue, mergeCombiners, numPartitions)

        return self.transform(func)

    def partitionBy(
        self: "DStream[Tuple[K, V]]",
        numPartitions: int,
        partitionFunc: Callable[[K], int] = portable_hash,
    ) -> "DStream[Tuple[K, V]]":
        """
        Return a copy of the DStream in which each RDD are partitioned
        using the specified partitioner.
        """
        return self.transform(lambda rdd: rdd.partitionBy(numPartitions, partitionFunc))

    @overload
    def foreachRDD(self: "DStream[T]", func: Callable[[RDD[T]], None]) -> None:
        ...

    @overload
    def foreachRDD(self: "DStream[T]", func: Callable[[datetime, RDD[T]], None]) -> None:
        ...

    def foreachRDD(
        self: "DStream[T]",
        func: Union[Callable[[RDD[T]], None], Callable[[datetime, RDD[T]], None]],
    ) -> None:
        """
        Apply a function to each RDD in this DStream.
        """
        if func.__code__.co_argcount == 1:
            old_func = func

            def func(_: datetime, rdd: "RDD[T]") -> None:
                return old_func(rdd)  # type: ignore[call-arg, arg-type]

        jfunc = TransformFunction(self._sc, func, self._jrdd_deserializer)
        assert self._ssc._jvm is not None
        api = self._ssc._jvm.PythonDStream
        api.callForeachRDD(self._jdstream, jfunc)

    def pprint(self, num: int = 10) -> None:
        """
        Print the first num elements of each RDD generated in this DStream.

        Parameters
        ----------
        num : int, optional
            the number of elements from the first will be printed.
        """

        def takeAndPrint(time: datetime, rdd: RDD[T]) -> None:
            taken = rdd.take(num + 1)
            print("-------------------------------------------")
            print("Time: %s" % time)
            print("-------------------------------------------")
            for record in taken[:num]:
                print(record)
            if len(taken) > num:
                print("...")
            print("")

        self.foreachRDD(takeAndPrint)

    def mapValues(self: "DStream[Tuple[K, V]]", f: Callable[[V], U]) -> "DStream[Tuple[K, U]]":
        """
        Return a new DStream by applying a map function to the value of
        each key-value pairs in this DStream without changing the key.
        """

        def map_values_fn(kv: Tuple[K, V]) -> Tuple[K, U]:
            return kv[0], f(kv[1])

        return self.map(map_values_fn, preservesPartitioning=True)

    def flatMapValues(
        self: "DStream[Tuple[K, V]]", f: Callable[[V], Iterable[U]]
    ) -> "DStream[Tuple[K, U]]":
        """
        Return a new DStream by applying a flatmap function to the value
        of each key-value pairs in this DStream without changing the key.
        """

        def flat_map_fn(kv: Tuple[K, V]) -> Iterable[Tuple[K, U]]:
            return ((kv[0], x) for x in f(kv[1]))

        return self.flatMap(flat_map_fn, preservesPartitioning=True)

    def glom(self: "DStream[T]") -> "DStream[List[T]]":
        """
        Return a new DStream in which RDD is generated by applying glom()
        to RDD of this DStream.
        """

        def func(iterator: Iterable[T]) -> Iterable[List[T]]:
            yield list(iterator)

        return self.mapPartitions(func)

    def cache(self: "DStream[T]") -> "DStream[T]":
        """
        Persist the RDDs of this DStream with the default storage level
        (`MEMORY_ONLY`).
        """
        self.is_cached = True
        self.persist(StorageLevel.MEMORY_ONLY)
        return self

    def persist(self: "DStream[T]", storageLevel: StorageLevel) -> "DStream[T]":
        """
        Persist the RDDs of this DStream with the given storage level
        """
        self.is_cached = True
        javaStorageLevel = self._sc._getJavaStorageLevel(storageLevel)
        self._jdstream.persist(javaStorageLevel)
        return self

    def checkpoint(self: "DStream[T]", interval: int) -> "DStream[T]":
        """
        Enable periodic checkpointing of RDDs of this DStream

        Parameters
        ----------
        interval : int
            time in seconds, after each period of that, generated
            RDD will be checkpointed
        """
        self.is_checkpointed = True
        self._jdstream.checkpoint(self._ssc._jduration(interval))
        return self

    def groupByKey(
        self: "DStream[Tuple[K, V]]", numPartitions: Optional[int] = None
    ) -> "DStream[Tuple[K, Iterable[V]]]":
        """
        Return a new DStream by applying groupByKey on each RDD.
        """
        if numPartitions is None:
            numPartitions = self._sc.defaultParallelism
        return self.transform(lambda rdd: rdd.groupByKey(numPartitions))

    def countByValue(self: "DStream[K]") -> "DStream[Tuple[K, int]]":
        """
        Return a new DStream in which each RDD contains the counts of each
        distinct value in each RDD of this DStream.
        """
        return self.map(lambda x: (x, 1)).reduceByKey(lambda x, y: x + y)

    def saveAsTextFiles(self, prefix: str, suffix: Optional[str] = None) -> None:
        """
        Save each RDD in this DStream as at text file, using string
        representation of elements.
        """

        def saveAsTextFile(t: Optional[datetime], rdd: RDD[T]) -> None:
            path = rddToFileName(prefix, suffix, t)
            try:
                rdd.saveAsTextFile(path)
            except Py4JJavaError as e:
                # after recovered from checkpointing, the foreachRDD may
                # be called twice
                if "FileAlreadyExistsException" not in str(e):
                    raise

        return self.foreachRDD(saveAsTextFile)

    # TODO: uncomment this until we have ssc.pickleFileStream()
    # def saveAsPickleFiles(self, prefix, suffix=None):
    #     """
    #     Save each RDD in this DStream as at binary file, the elements are
    #     serialized by pickle.
    #     """
    #     def saveAsPickleFile(t, rdd):
    #         path = rddToFileName(prefix, suffix, t)
    #         try:
    #             rdd.saveAsPickleFile(path)
    #         except Py4JJavaError as e:
    #             # after recovered from checkpointing, the foreachRDD may
    #             # be called twice
    #             if 'FileAlreadyExistsException' not in str(e):
    #                 raise
    #     return self.foreachRDD(saveAsPickleFile)

    @overload
    def transform(self: "DStream[T]", func: Callable[[RDD[T]], RDD[U]]) -> "TransformedDStream[U]":
        ...

    @overload
    def transform(
        self: "DStream[T]", func: Callable[[datetime, RDD[T]], RDD[U]]
    ) -> "TransformedDStream[U]":
        ...

    def transform(
        self: "DStream[T]",
        func: Union[Callable[[RDD[T]], RDD[U]], Callable[[datetime, RDD[T]], RDD[U]]],
    ) -> "TransformedDStream[U]":
        """
        Return a new DStream in which each RDD is generated by applying a function
        on each RDD of this DStream.

        `func` can have one argument of `rdd`, or have two arguments of
        (`time`, `rdd`)
        """
        if func.__code__.co_argcount == 1:
            oldfunc = func

            def func(_: datetime, rdd: RDD[T]) -> RDD[U]:
                return oldfunc(rdd)  # type: ignore[arg-type, call-arg]

        assert func.__code__.co_argcount == 2, "func should take one or two arguments"
        return TransformedDStream(self, func)

    @overload
    def transformWith(
        self: "DStream[T]",
        func: Callable[[RDD[T], RDD[U]], RDD[V]],
        other: "DStream[U]",
        keepSerializer: bool = ...,
    ) -> "DStream[V]":
        ...

    @overload
    def transformWith(
        self: "DStream[T]",
        func: Callable[[datetime, RDD[T], RDD[U]], RDD[V]],
        other: "DStream[U]",
        keepSerializer: bool = ...,
    ) -> "DStream[V]":
        ...

    def transformWith(
        self: "DStream[T]",
        func: Union[
            Callable[[RDD[T], RDD[U]], RDD[V]],
            Callable[[datetime, RDD[T], RDD[U]], RDD[V]],
        ],
        other: "DStream[U]",
        keepSerializer: bool = False,
    ) -> "DStream[V]":
        """
        Return a new DStream in which each RDD is generated by applying a function
        on each RDD of this DStream and 'other' DStream.

        `func` can have two arguments of (`rdd_a`, `rdd_b`) or have three
        arguments of (`time`, `rdd_a`, `rdd_b`)
        """
        if func.__code__.co_argcount == 2:
            oldfunc = func

            def func(_: datetime, a: RDD[T], b: RDD[U]) -> RDD[V]:
                return oldfunc(a, b)  # type: ignore[call-arg, arg-type]

        assert func.__code__.co_argcount == 3, "func should take two or three arguments"
        jfunc = TransformFunction(
            self._sc,
            func,
            self._jrdd_deserializer,
            other._jrdd_deserializer,
        )
        assert self._sc._jvm is not None
        dstream = self._sc._jvm.PythonTransformed2DStream(
            self._jdstream.dstream(), other._jdstream.dstream(), jfunc
        )
        jrdd_serializer = self._jrdd_deserializer if keepSerializer else self._sc.serializer
        return DStream(dstream.asJavaDStream(), self._ssc, jrdd_serializer)

    def repartition(self: "DStream[T]", numPartitions: int) -> "DStream[T]":
        """
        Return a new DStream with an increased or decreased level of parallelism.
        """
        return self.transform(lambda rdd: rdd.repartition(numPartitions))

    @property
    def _slideDuration(self) -> None:
        """
        Return the slideDuration in seconds of this DStream
        """
        return self._jdstream.dstream().slideDuration().milliseconds() / 1000.0

    def union(self: "DStream[T]", other: "DStream[U]") -> "DStream[Union[T, U]]":
        """
        Return a new DStream by unifying data of another DStream with this DStream.

        Parameters
        ----------
        other : :class:`DStream`
            Another DStream having the same interval (i.e., slideDuration)
            as this DStream.
        """
        if self._slideDuration != other._slideDuration:
            raise ValueError("the two DStream should have same slide duration")
        return self.transformWith(lambda a, b: a.union(b), other, True)

    def cogroup(
        self: "DStream[Tuple[K, V]]",
        other: "DStream[Tuple[K, U]]",
        numPartitions: Optional[int] = None,
    ) -> "DStream[Tuple[K, Tuple[ResultIterable[V], ResultIterable[U]]]]":
        """
        Return a new DStream by applying 'cogroup' between RDDs of this
        DStream and `other` DStream.

        Hash partitioning is used to generate the RDDs with `numPartitions` partitions.
        """
        if numPartitions is None:
            numPartitions = self._sc.defaultParallelism
        return self.transformWith(
            lambda a, b: a.cogroup(b, numPartitions),
            other,
        )

    def join(
        self: "DStream[Tuple[K, V]]",
        other: "DStream[Tuple[K, U]]",
        numPartitions: Optional[int] = None,
    ) -> "DStream[Tuple[K, Tuple[V, U]]]":
        """
        Return a new DStream by applying 'join' between RDDs of this DStream and
        `other` DStream.

        Hash partitioning is used to generate the RDDs with `numPartitions`
        partitions.
        """
        if numPartitions is None:
            numPartitions = self._sc.defaultParallelism
        return self.transformWith(lambda a, b: a.join(b, numPartitions), other)

    def leftOuterJoin(
        self: "DStream[Tuple[K, V]]",
        other: "DStream[Tuple[K, U]]",
        numPartitions: Optional[int] = None,
    ) -> "DStream[Tuple[K, Tuple[V, Optional[U]]]]":
        """
        Return a new DStream by applying 'left outer join' between RDDs of this DStream and
        `other` DStream.

        Hash partitioning is used to generate the RDDs with `numPartitions`
        partitions.
        """
        if numPartitions is None:
            numPartitions = self._sc.defaultParallelism
        return self.transformWith(lambda a, b: a.leftOuterJoin(b, numPartitions), other)

    def rightOuterJoin(
        self: "DStream[Tuple[K, V]]",
        other: "DStream[Tuple[K, U]]",
        numPartitions: Optional[int] = None,
    ) -> "DStream[Tuple[K, Tuple[Optional[V], U]]]":
        """
        Return a new DStream by applying 'right outer join' between RDDs of this DStream and
        `other` DStream.

        Hash partitioning is used to generate the RDDs with `numPartitions`
        partitions.
        """
        if numPartitions is None:
            numPartitions = self._sc.defaultParallelism
        return self.transformWith(lambda a, b: a.rightOuterJoin(b, numPartitions), other)

    def fullOuterJoin(
        self: "DStream[Tuple[K, V]]",
        other: "DStream[Tuple[K, U]]",
        numPartitions: Optional[int] = None,
    ) -> "DStream[Tuple[K, Tuple[Optional[V], Optional[U]]]]":
        """
        Return a new DStream by applying 'full outer join' between RDDs of this DStream and
        `other` DStream.

        Hash partitioning is used to generate the RDDs with `numPartitions`
        partitions.
        """
        if numPartitions is None:
            numPartitions = self._sc.defaultParallelism
        return self.transformWith(lambda a, b: a.fullOuterJoin(b, numPartitions), other)

    def _jtime(self, timestamp: Union[datetime, int, float]) -> JavaObject:
        """Convert datetime or unix_timestamp into Time"""
        if isinstance(timestamp, datetime):
            timestamp = time.mktime(timestamp.timetuple())
        assert self._sc._jvm is not None
        return self._sc._jvm.Time(int(timestamp * 1000))

    def slice(self, begin: Union[datetime, int], end: Union[datetime, int]) -> List[RDD[T]]:
        """
        Return all the RDDs between 'begin' to 'end' (both included)

        `begin`, `end` could be datetime.datetime() or unix_timestamp
        """
        jrdds = self._jdstream.slice(self._jtime(begin), self._jtime(end))
        return [RDD(jrdd, self._sc, self._jrdd_deserializer) for jrdd in jrdds]

    def _validate_window_param(self, window: int, slide: Optional[int]) -> None:
        duration = self._jdstream.dstream().slideDuration().milliseconds()
        if int(window * 1000) % duration != 0:
            raise ValueError(
                "windowDuration must be multiple of the parent "
                "dstream's slide (batch) duration (%d ms)" % duration
            )
        if slide and int(slide * 1000) % duration != 0:
            raise ValueError(
                "slideDuration must be multiple of the parent "
                "dstream's slide (batch) duration (%d ms)" % duration
            )

    def window(self, windowDuration: int, slideDuration: Optional[int] = None) -> "DStream[T]":
        """
        Return a new DStream in which each RDD contains all the elements in seen in a
        sliding window of time over this DStream.

        Parameters
        ----------
        windowDuration : int
            width of the window; must be a multiple of this DStream's
            batching interval
        slideDuration : int, optional
            sliding interval of the window (i.e., the interval after which
            the new DStream will generate RDDs); must be a multiple of this
            DStream's batching interval
        """
        self._validate_window_param(windowDuration, slideDuration)
        d = self._ssc._jduration(windowDuration)
        if slideDuration is None:
            return DStream(self._jdstream.window(d), self._ssc, self._jrdd_deserializer)
        s = self._ssc._jduration(slideDuration)
        return DStream(self._jdstream.window(d, s), self._ssc, self._jrdd_deserializer)

    def reduceByWindow(
        self: "DStream[T]",
        reduceFunc: Callable[[T, T], T],
        invReduceFunc: Optional[Callable[[T, T], T]],
        windowDuration: int,
        slideDuration: int,
    ) -> "DStream[T]":
        """
        Return a new DStream in which each RDD has a single element generated by reducing all
        elements in a sliding window over this DStream.

        if `invReduceFunc` is not None, the reduction is done incrementally
        using the old window's reduced value :

        1. reduce the new values that entered the window (e.g., adding new counts)

        2. "inverse reduce" the old values that left the window (e.g., subtracting old counts)
        This is more efficient than `invReduceFunc` is None.

        Parameters
        ----------
        reduceFunc : function
            associative and commutative reduce function
        invReduceFunc : function
            inverse reduce function of `reduceFunc`; such that for all y,
            and invertible x:
            `invReduceFunc(reduceFunc(x, y), x) = y`
        windowDuration : int
            width of the window; must be a multiple of this DStream's
            batching interval
        slideDuration : int
            sliding interval of the window (i.e., the interval after which
            the new DStream will generate RDDs); must be a multiple of this
            DStream's batching interval
        """
        keyed = self.map(lambda x: (1, x))
        reduced = keyed.reduceByKeyAndWindow(
            reduceFunc, invReduceFunc, windowDuration, slideDuration, 1
        )
        return reduced.map(lambda kv: kv[1])

    def countByWindow(
        self: "DStream[T]", windowDuration: int, slideDuration: int
    ) -> "DStream[int]":
        """
        Return a new DStream in which each RDD has a single element generated
        by counting the number of elements in a window over this DStream.
        windowDuration and slideDuration are as defined in the window() operation.

        This is equivalent to window(windowDuration, slideDuration).count(),
        but will be more efficient if window is large.
        """
        return self.map(lambda x: 1).reduceByWindow(
            operator.add, operator.sub, windowDuration, slideDuration
        )

    def countByValueAndWindow(
        self: "DStream[T]",
        windowDuration: int,
        slideDuration: int,
        numPartitions: Optional[int] = None,
    ) -> "DStream[Tuple[T, int]]":
        """
        Return a new DStream in which each RDD contains the count of distinct elements in
        RDDs in a sliding window over this DStream.

        Parameters
        ----------
        windowDuration : int
            width of the window; must be a multiple of this DStream's
            batching interval
        slideDuration : int
            sliding interval of the window (i.e., the interval after which
            the new DStream will generate RDDs); must be a multiple of this
            DStream's batching interval
        numPartitions : int, optional
            number of partitions of each RDD in the new DStream.
        """
        keyed = self.map(lambda x: (x, 1))
        counted = keyed.reduceByKeyAndWindow(
            operator.add, operator.sub, windowDuration, slideDuration, numPartitions
        )
        return counted.filter(lambda kv: kv[1] > 0)

    def groupByKeyAndWindow(
        self: "DStream[Tuple[K, V]]",
        windowDuration: int,
        slideDuration: int,
        numPartitions: Optional[int] = None,
    ) -> "DStream[Tuple[K, Iterable[V]]]":
        """
        Return a new DStream by applying `groupByKey` over a sliding window.
        Similar to `DStream.groupByKey()`, but applies it over a sliding window.

        Parameters
        ----------
        windowDuration : int
            width of the window; must be a multiple of this DStream's
            batching interval
        slideDuration : int
            sliding interval of the window (i.e., the interval after which
            the new DStream will generate RDDs); must be a multiple of this
            DStream's batching interval
        numPartitions : int, optional
            Number of partitions of each RDD in the new DStream.
        """
        ls = self.mapValues(lambda x: [x])
        grouped = ls.reduceByKeyAndWindow(
            lambda a, b: a.extend(b) or a,  # type: ignore[func-returns-value]
            lambda a, b: a[len(b) :],
            windowDuration,
            slideDuration,
            numPartitions,
        )
        return grouped.mapValues(ResultIterable)

    def reduceByKeyAndWindow(
        self: "DStream[Tuple[K, V]]",
        func: Callable[[V, V], V],
        invFunc: Optional[Callable[[V, V], V]],
        windowDuration: int,
        slideDuration: Optional[int] = None,
        numPartitions: Optional[int] = None,
        filterFunc: Optional[Callable[[Tuple[K, V]], bool]] = None,
    ) -> "DStream[Tuple[K, V]]":
        """
        Return a new DStream by applying incremental `reduceByKey` over a sliding window.

        The reduced value of over a new window is calculated using the old window's reduce value :
         1. reduce the new values that entered the window (e.g., adding new counts)
         2. "inverse reduce" the old values that left the window (e.g., subtracting old counts)

        `invFunc` can be None, then it will reduce all the RDDs in window, could be slower
        than having `invFunc`.

        Parameters
        ----------
        func : function
            associative and commutative reduce function
        invFunc : function
            inverse function of `reduceFunc`
        windowDuration : int
            width of the window; must be a multiple of this DStream's
            batching interval
        slideDuration : int, optional
            sliding interval of the window (i.e., the interval after which
            the new DStream will generate RDDs); must be a multiple of this
            DStream's batching interval
        numPartitions : int, optional
            number of partitions of each RDD in the new DStream.
        filterFunc : function, optional
            function to filter expired key-value pairs;
            only pairs that satisfy the function are retained
            set this to null if you do not want to filter
        """
        self._validate_window_param(windowDuration, slideDuration)
        if numPartitions is None:
            numPartitions = self._sc.defaultParallelism

        reduced = self.reduceByKey(func, numPartitions)

        if invFunc:

            def reduceFunc(t: datetime, a: Any, b: Any) -> Any:
                b = b.reduceByKey(func, numPartitions)
                r = a.union(b).reduceByKey(func, numPartitions) if a else b
                if filterFunc:
                    r = r.filter(filterFunc)
                return r

            def invReduceFunc(t: datetime, a: Any, b: Any) -> Any:
                b = b.reduceByKey(func, numPartitions)
                joined = a.leftOuterJoin(b, numPartitions)
                return joined.mapValues(
                    lambda kv: invFunc(kv[0], kv[1]) if kv[1] is not None else kv[0]
                )

            jreduceFunc = TransformFunction(self._sc, reduceFunc, reduced._jrdd_deserializer)
            jinvReduceFunc = TransformFunction(self._sc, invReduceFunc, reduced._jrdd_deserializer)
            if slideDuration is None:
                slideDuration = self._slideDuration
            assert self._sc._jvm is not None
            dstream = self._sc._jvm.PythonReducedWindowedDStream(
                reduced._jdstream.dstream(),
                jreduceFunc,
                jinvReduceFunc,
                self._ssc._jduration(windowDuration),
                self._ssc._jduration(slideDuration),  # type: ignore[arg-type]
            )
            return DStream(dstream.asJavaDStream(), self._ssc, self._sc.serializer)
        else:
            return reduced.window(windowDuration, slideDuration).reduceByKey(
                func, numPartitions  # type: ignore[arg-type]
            )

    def updateStateByKey(
        self: "DStream[Tuple[K, V]]",
        updateFunc: Callable[[Iterable[V], Optional[S]], S],
        numPartitions: Optional[int] = None,
        initialRDD: Optional[Union[RDD[Tuple[K, S]], Iterable[Tuple[K, S]]]] = None,
    ) -> "DStream[Tuple[K, S]]":
        """
        Return a new "state" DStream where the state for each key is updated by applying
        the given function on the previous state of the key and the new values of the key.

        Parameters
        ----------
        updateFunc : function
            State update function. If this function returns None, then
            corresponding state key-value pair will be eliminated.
        """
        if numPartitions is None:
            numPartitions = self._sc.defaultParallelism

        if initialRDD and not isinstance(initialRDD, RDD):
            initialRDD = self._sc.parallelize(initialRDD)

        def reduceFunc(t: datetime, a: Any, b: Any) -> Any:
            if a is None:
                g = b.groupByKey(numPartitions).mapValues(lambda vs: (list(vs), None))
            else:
                g = a.cogroup(b.partitionBy(numPartitions), numPartitions)
                g = g.mapValues(lambda ab: (list(ab[1]), list(ab[0])[0] if len(ab[0]) else None))
            state = g.mapValues(lambda vs_s: updateFunc(vs_s[0], vs_s[1]))
            return state.filter(lambda k_v: k_v[1] is not None)

        jreduceFunc = TransformFunction(
            self._sc,
            reduceFunc,
            self._sc.serializer,
            self._jrdd_deserializer,
        )
        if initialRDD:
            initialRDD = cast(RDD[Tuple[K, S]], initialRDD)._reserialize(self._jrdd_deserializer)
            assert self._sc._jvm is not None
            dstream = self._sc._jvm.PythonStateDStream(
                self._jdstream.dstream(),
                jreduceFunc,
                initialRDD._jrdd,
            )
        else:
            assert self._sc._jvm is not None
            dstream = self._sc._jvm.PythonStateDStream(self._jdstream.dstream(), jreduceFunc)

        return DStream(dstream.asJavaDStream(), self._ssc, self._sc.serializer)


class TransformedDStream(DStream[U]):
    """
    TransformedDStream is a DStream generated by an Python function
    transforming each RDD of a DStream to another RDDs.

    Multiple continuous transformations of DStream can be combined into
    one transformation.
    """

    @overload
    def __init__(self: DStream[U], prev: DStream[T], func: Callable[[RDD[T]], RDD[U]]):
        ...

    @overload
    def __init__(
        self: DStream[U],
        prev: DStream[T],
        func: Callable[[datetime, RDD[T]], RDD[U]],
    ):
        ...

    def __init__(
        self,
        prev: DStream[T],
        func: Union[Callable[[RDD[T]], RDD[U]], Callable[[datetime, RDD[T]], RDD[U]]],
    ):
        self._ssc = prev._ssc
        self._sc = self._ssc._sc
        self._jrdd_deserializer = self._sc.serializer
        self.is_cached = False
        self.is_checkpointed = False
        self._jdstream_val = None

        # Using type() to avoid folding the functions and compacting the DStreams which is not
        # not strictly an object of TransformedDStream.
        if type(prev) is TransformedDStream and not prev.is_cached and not prev.is_checkpointed:
            prev_func: Callable = prev.func
            func = cast(Callable[[datetime, RDD[T]], RDD[U]], func)
            self.func: Union[
                Callable[[RDD[T]], RDD[U]], Callable[[datetime, RDD[T]], RDD[U]]
            ] = lambda t, rdd: func(t, prev_func(t, rdd))
            self.prev: DStream[T] = prev.prev
        else:
            self.prev = prev
            self.func = func

    @property
    def _jdstream(self) -> JavaObject:
        if self._jdstream_val is not None:
            return self._jdstream_val

        jfunc = TransformFunction(self._sc, self.func, self.prev._jrdd_deserializer)
        assert self._sc._jvm is not None
        dstream = self._sc._jvm.PythonTransformedDStream(self.prev._jdstream.dstream(), jfunc)
        self._jdstream_val = dstream.asJavaDStream()
        return self._jdstream_val