from typing import Generic, TypeVar

from erdos.config import OperatorConfig
from erdos.internal import PyTimestamp
from erdos.streams import WriteStream
from erdos.timestamp import Timestamp

T = TypeVar("T")
U = TypeVar("U")
V = TypeVar("V")


class SinkContext:
    """A :py:class:`SinkContext` instance enables developers to retrieve
    metadata about the current invocation of either a message or a watermark
    callback in a :py:class:`.Sink` operator.

    Attributes:
        timestamp (:py:class:`.Timestamp`): The timestamp of the current
            invocation of the callback.
        config (:py:class:`.OperatorConfig`): The operator config generated by
            the driver upon connection of the operator to the graph.
    """

    def __init__(self, timestamp: PyTimestamp, config: OperatorConfig):
        self.timestamp = Timestamp(_py_timestamp=timestamp)
        self.config = config

    def __str__(self) -> str:
        return "SinkContext(Timestamp={}, Config={})".format(
            self.timestamp, self.config
        )


class OneInOneOutContext(Generic[T]):
    """A :py:class:`OneInOneOutContext` instance enables developers to retrieve
    metadata about the current invocation of either a message or a watermark
    callback in a :py:class:`.OneInOneOut` operator.

    Attributes:
        timestamp (:py:class:`.Timestamp`): The timestamp of the current
            invocation of the callback.
        config (:py:class:`.OperatorConfig`): The operator config generated by
            the driver upon connection of the operator to the graph.
        write_stream (:py:class:`.WriteStream`): The write stream to send
            results to downstream operators.
    """

    def __init__(
        self,
        timestamp: PyTimestamp,
        config: OperatorConfig,
        write_stream: WriteStream[T],
    ):
        self.timestamp = Timestamp(_py_timestamp=timestamp)
        self.config = config
        self.write_stream = write_stream

    def __str__(self) -> str:
        return "OneInOneOutContext(Timestamp={}, Config={}, WriteStream={})".format(
            self.timestamp, self.config, self.write_stream.name
        )


class OneInTwoOutContext(Generic[T, U]):
    """A :py:class:`OneInTwoOutContext` instance enables developers to retrieve
    metadata about the current invocation of either a message or a watermark
    callback in a :py:class:`.OneInTwoOut` operator.

    Attributes:
        timestamp (:py:class:`.Timestamp`): The timestamp of the current
            invocation of the callback.
        config (:py:class:`.OperatorConfig`): The operator config generated
            by the driver upon connection of the operator to the graph.
        left_write_stream (:py:class:`.WriteStream`): The first write stream to
            send results to downstream operators.
        right_write_stream (:py:class:`.WriteStream`): The second write stream
            to send results to downstream operators.
    """

    def __init__(
        self,
        timestamp: PyTimestamp,
        config: OperatorConfig,
        left_write_stream: WriteStream[T],
        right_write_stream: WriteStream[U],
    ):
        self.timestamp = Timestamp(_py_timestamp=timestamp)
        self.config = config
        self.left_write_stream = left_write_stream
        self.right_write_stream = right_write_stream

    def __str__(self) -> str:
        return "OneInTwoOutContext(Timestamp={}, Config={}, \
                Left WriteStream={}, Right WriteStream={})".format(
            self.timestamp,
            self.config,
            self.left_write_stream.name,
            self.right_write_stream.name,
        )


class TwoInOneOutContext(Generic[T]):
    """A :py:class:`TwoInOneOutContext` instance enables developers to retrieve
    metadata about the current invocation of either a message or a watermark
    callback in a :py:class:`.TwoInOneOut` operator.

    Attributes:
        timestamp (:py:class:`.Timestamp`): The timestamp of the current
            invocation of the callback.
        config (:py:class:`.OperatorConfig`): The operator config generated by
            the driver upon connection of the operator to the graph.
        write_stream (:py:class:`.WriteStream`): The write stream to send
            results to downstream operators.
    """

    def __init__(
        self,
        timestamp: PyTimestamp,
        config: OperatorConfig,
        write_stream: WriteStream[T],
    ):
        self.timestamp = Timestamp(_py_timestamp=timestamp)
        self.config = config
        self.write_stream = write_stream

    def __str__(self) -> str:
        return "TwoInOneOutContext(Timestamp={}, Config={}, WriteStream={})".format(
            self.timestamp,
            self.config,
            self.write_stream.name,
        )
