import copy
import heapq
import logging
import pickle
import sys
import time
from enum import Enum

from erdos.data_stream import DataStream
from erdos.message import Message
from erdos.message import WatermarkMessage
from erdos.op import Op
from erdos.timestamp import Timestamp
from erdos.utils import frequency


class LogOp(Op):
    """Operator which logs data from input streams to file or stdout.

    Args:
        input_streams (list): list of input streams from which to log.
        name (str): unique name for this operator. Generated by default.
        fmt (str): formatting string for logged messages. See logging module
            documentation for more details.
        date_fmt (str): formatting string for the date. See logging module
            documentation for more details.
        console_output (bool): whether to print messages to stdout.
        filename (str): file at which to log messages. If empty, messages are
            not logged to file.
        mode (str): mode in which the log file is opened.
        encoding (str): encoding with which the log file is opened.
    """

    def __init__(self,
                 name="",
                 fmt=None,
                 date_fmt=None,
                 console_output=True,
                 filename="",
                 mode="a",
                 encoding=None):
        super(LogOp, self).__init__(name)
        self.format = fmt
        self.date_format = date_fmt
        self.console_output = console_output
        self.filename = filename
        self.mode = mode
        self.encoding = encoding

    @staticmethod
    def setup_streams(input_streams):
        input_streams.add_callback(LogOp.log_input)
        return []

    def log_input(self, msg):
        self.logger.info(msg)

    def execute(self):
        # Set up logger
        self.logger = logging.getLogger(self.name)
        self.logger.setLevel(logging.INFO)

        formatter = logging.Formatter(self.format, self.date_format)

        if self.console_output:
            stream_handler = logging.StreamHandler(sys.stdout)
            stream_handler.setFormatter(formatter)
            self.logger.addHandler(stream_handler)

        if self.filename:
            file_handler = logging.FileHandler(self.filename, self.mode,
                                               self.encoding)
            file_handler.setFormatter(formatter)
            self.logger.addHandler(file_handler)
        self.logger.propagate = False
        self.spin()


class RecordOp(Op):
    """Operator which saves serialized data from input streams to file.

    Args:
        filename (str): path to file.
        input_streams (list): list of input streams from which to save data.
        name (str): unique name for this operator. Generated by default.
    """

    def __init__(self, name, filename):
        super(RecordOp, self).__init__(name)
        self.filename = filename

    @staticmethod
    def setup_streams(input_streams, filter):
        if filter:
            input_streams = input_streams.filter_name(filter)
        input_streams.add_callback(RecordOp.record_data)
        return []

    def record_data(self, msg):
        pickle.dump(msg, self._file)

    def execute(self):
        self._file = open(self.filename, "wb")

        # Write input stream info
        input_stream_info = [(input_stream.data_type, input_stream.name)
                             for input_stream in self.input_streams]
        pickle.dump(input_stream_info, self._file)
        self.spin()


class ReplayOp(Op):
    """Operator which replays saved data from file to an output stream.

    Args:
        filename (str): path to file.
        frequency (int): rate at which the operator publishes data. If 0, the
            operator publishes data as soon as it is read from file.
        name (str): unique name for this operator. Generated by default.
    """

    def __init__(self, filename, frequency=0, name="replay_op"):
        super(ReplayOp, self).__init__(name)
        self.filename = filename
        self.frequency = frequency

    @staticmethod
    def setup_streams(input_streams, filename):
        output_streams = []
        with open(filename, "rb") as f:
            stream_info = pickle.load(f)
            for data_type, name in stream_info:
                output_streams.append(
                    DataStream(data_type=data_type, name=name))
        return output_streams

    def publish_data(self):
        if self._file.closed:
            return
        try:
            msg = pickle.load(self._file)
            self.get_output_stream(msg.stream_name).send(msg)
        except EOFError:
            logging.error("Reached end of input file: {0}".format(
                self.filename))
            self._file.close()

    def execute(self):
        self._file = open(self.filename, "rb")
        pickle.load(self._file)  # Read past the stream names

        if self.frequency:
            publish_data = frequency(self.frequency)(self.publish_data)
            publish_data()
            self.spin()
        else:
            while True:
                self.publish_data()
        self.spin()


class FileWriterOp(Op):
    def __init__(self, name, file_name):
        super(FileWriterOp, self).__init__(name)
        self._output_file = open(file_name, 'w')

    def __del__(self):
        self._output_file.close()

    @staticmethod
    def setup_streams(input_streams, filter_stream_lambda=None):
        input_streams.filter(filter_stream_lambda).add_callback(
            FileWriterOp.on_msg)
        return []

    def on_msg(self, msg):
        self._output_file.write(str(msg.data) + '\n')


class WhereOp(Op):
    def __init__(self, name, output_stream_name, where_lambda):
        super(WhereOp, self).__init__(name)
        self._output_stream_name = output_stream_name
        self._where_lambda = where_lambda

    @staticmethod
    def setup_streams(input_streams,
                      output_stream_name,
                      filter_stream_lambda=None):
        input_streams.filter(filter_stream_lambda).add_callback(WhereOp.on_msg)
        return [DataStream(name=output_stream_name)]

    def on_msg(self, msg):
        if self._where_lambda(msg):
            self.get_output_stream(self._output_stream_name).send(msg)


class MapOp(Op):
    def __init__(self, name, output_stream_name, map_lambda):
        super(MapOp, self).__init__(name)
        self._output_stream_name = output_stream_name
        self._map_lambda = map_lambda

    @staticmethod
    def setup_streams(input_streams,
                      output_stream_name,
                      filter_stream_lambda=None):
        input_streams.filter(filter_stream_lambda).add_callback(MapOp.on_msg)
        return [DataStream(name=output_stream_name)]

    def on_msg(self, msg):
        new_data = self._map_lambda(msg)
        new_msg = Message(new_data, msg.timestamp)
        self.get_output_stream(self._output_stream_name).send(new_msg)


class MapManyOp(Op):
    def __init__(self, name, output_stream_name, map_lambda=None):
        super(MapManyOp, self).__init__(name)
        self._output_stream_name = output_stream_name
        self._map_lambda = map_lambda

    @staticmethod
    def setup_streams(input_streams,
                      output_stream_name,
                      filter_stream_lambda=None):
        input_streams.filter(filter_stream_lambda).add_callback(
            MapManyOp.on_msg)
        return [DataStream(name=output_stream_name)]

    def on_msg(self, msg):
        if self._map_lambda:
            data = self._map_lambda(msg)
        else:
            data = msg.data
        try:
            data = iter(msg.data)
        except TypeError:
            data = [data]
        for val in data:
            self.get_output_stream(self._output_stream_name).send(
                Message(val, msg.timestamp))


class ConcatOp(Op):
    def __init__(self, name, output_stream_name):
        super(ConcatOp, self).__init__(name)
        self._output_stream_name = output_stream_name

    @staticmethod
    def setup_streams(input_streams,
                      output_stream_name,
                      filter_stream_lambda=None):
        input_streams.filter(filter_stream_lambda).add_callback(
            ConcatOp.on_msg)
        return [DataStream(name=output_stream_name)]

    def on_msg(self, msg):
        self.get_output_stream(self._output_stream_name).send(msg)


class UnzipOp(Op):
    def __init__(self, name, output_stream_name1, output_stream_name2):
        super(UnzipOp, self).__init__(name)
        self._output_stream_name1 = output_stream_name1
        self._output_stream_name2 = output_stream_name2

    @staticmethod
    def setup_streams(input_streams,
                      output_stream_name1,
                      output_stream_name2,
                      filter_stream_lambda=None):
        input_streams.filter(filter_stream_lambda).add_callback(UnzipOp.on_msg)
        return [
            DataStream(name=output_stream_name1),
            DataStream(name=output_stream_name2)
        ]

    def on_msg(self, msg):
        (left_val, right_val) = msg.data
        self.get_output_stream(self._output_stream_name1).send(
            Message(left_val, Timestamp(msg.timestamp)))
        self.get_output_stream(self._output_stream_name2).send(
            Message(right_val, Timestamp(msg.timestamp)))


def get_timestamp_ms():
    return int(time.time() * 1000)


def get_window_start_with_offset(window_size_ms, window_offset_ms):
    return ((get_timestamp_ms() - window_offset_ms) // window_size_ms *
            window_size_ms)


class Window(object):
    """Window base class.

    Attributes:
        window_id(int): A unique int identifying the window.
    """

    def __init__(self, window_id, is_time_window=False):
        self.is_time_window = is_time_window
        self._uid = window_id

    @property
    def uid(self):
        return self._uid


class TimeWindow(Window):
    """Class for time-based windows.

    Attributes:
        start_time (int): Window starting time in ms
        end_time (int): Window ending time in ms
    """

    def __init__(self, start_time, end_time):
        super(TimeWindow, self).__init__(hash((start_time, end_time)), True)
        self._start_time = start_time
        self._end_time = end_time

    @property
    def end_time(self):
        return self._end_time


class WindowAssigner(object):
    """WindowAssigner base class.
    """

    def assign_windows(self, msg):
        """Assigns a message to a list of windows."""
        raise NotImplementedError(
            'Window assigner must implement assign_windows')


class TumblingWindowAssigner(WindowAssigner):
    """Assigns messages to non-overlapping windows (i.e., tumbling).

    Attributes:
        time_size_ms (int): Time duration of the window
        time_offset_ms (int): Offset of the window
    """

    def __init__(self, time_size_ms, time_offset_ms=0):
        self._size = time_size_ms
        self._offset = time_offset_ms

    def assign_windows(self, msg):
        start_time = get_window_start_with_offset(self._size, self._offset)
        return [TimeWindow(start_time, start_time + self._size)]


class SlidingWindowAssigner(WindowAssigner):
    """Assigns messages to sliding windows.

    Attributes:
        time_size_ms (int): Time duration of the window
        time_slide_ms (int): By how much time a window slides
        time_offset_ms (int): Offset of the window
    """

    def __init__(self, time_size_ms, time_slide_ms, time_offset_ms=0):
        self._size = time_size_ms
        self._slide = time_slide_ms
        self._offset = time_offset_ms

    def assign_windows(self, msg):
        start_time = get_window_start_with_offset(self._size, self._offset)
        windows = []
        while start_time < msg.timestamp.coordinates[0]:
            windows.append(TimeWindow(start_time, start_time + self._size))
            start_time += self._slide
        return windows


class CountWindowAssigner(WindowAssigner):
    """Assigns n messages to each window.
    Attributes:
        msg_count (int): Number of messages per window.
    """

    def __init__(self, msg_count):
        self._msg_count = msg_count
        self._cur_msg_count = 0
        self._window_id = 0

    def assign_windows(self, msg):
        if self._cur_msg_count < self._msg_count:
            self._cur_msg_count += 1
        else:
            self._cur_msg_count = 1
            self._window_id += 1
        return [Window(self._window_id)]


class WindowTrigger(object):
    """WindowTrigger base class.

    All triggers must inherit from this class, and must implement:
    1. is_time_based: Return true if trigger is time-based.
    2. on_message: Invoked for each message received for the window.
    3. on_time: Invoked when the time triggers.
    """

    def is_time_based(self):
        """Returns True if the window is time-based."""
        raise NotImplementedError('Trigger must implement is_time_based')

    def on_message(self, msg, window):
        raise NotImplementedError('Trigger must implement on_message')

    def on_time(self, time, window):
        raise NotImplementedError('Trigger must implement on_time')


class WindowTriggerAction(Enum):
    CONTINUE = 1  # Continue accumulating messages, do no trigger processing.
    PROCESS = 2  # Trigger window processing.


class TimeWindowTrigger(WindowTrigger):
    """Time-based trigger."""

    def __init__(self):
        pass

    def is_time_based(self):
        return True

    def on_message(self, msg, window):
        return WindowTriggerAction.CONTINUE

    def on_time(self, time, window):
        return WindowTriggerAction.PROCESS


class CountWindowTrigger(WindowTrigger):
    """Count-based trigger.

    Attributes:
        msg_cnt (int): Number of messages after which to process the window.
    """

    def __init__(self, msg_cnt):
        self._msg_cnt = msg_cnt
        self._cur_count = 0

    def is_time_based(self):
        return False

    def on_message(self, msg, window):
        self._cur_count += 1
        if self._cur_count == self._msg_cnt:
            self._cur_count = 0
            return WindowTriggerAction.PROCESS
        return WindowTriggerAction.CONTINUE

    def on_time(self, time, window):
        return WindowTriggerAction.PROCESS


class WindowProcessor(object):
    """WindowProcessor base class.

    All window processors must inherit from this class, and must implement:
    1. process: Process messages within a window.
    """

    def __init__(self):
        pass

    def process(self, msgs):
        """Process messages within a window."""
        raise NotImplementedError('WindowProcessor must implement process')


class SumWindowProcessor(WindowProcessor):
    def process(self, msgs):
        assert len(msgs) > 0
        res = 0
        for msg in msgs:
            res += msg.data
        return [Message(res, msgs[0].timestamp)]


class WindowOp(Op):
    """ERDOS-provided operator for windowing streams.

    Attributes:
        name (str): A unique string naming the operator.
        output_stream_name (str): Name of the stream on which to output data.
        assigner (WindowAssigner): object that assigns messages to windows.
        trigger (WindowTrigger): object that decides when windows must be
            processed.
        processor (WindowProcessor): object used to process a window when it is
            ready.
    """

    def __init__(self, name, output_stream_name, assigner, trigger, processor):
        super(WindowOp, self).__init__(name)
        self._output_stream_name = output_stream_name
        self._assigner = assigner
        self._trigger = trigger
        self._processor = processor
        self._window_state_map = {}
        self._window_trigger_map = {}
        self._window_end_pqueue = []
        self._registered_windows = set([])

    @staticmethod
    def setup_streams(input_streams,
                      output_stream_name,
                      filter_stream_lambda=None):
        input_streams.filter(filter_stream_lambda).add_callback(
            WindowOp.on_msg)
        return [DataStream(name=output_stream_name)]

    def on_process(self, window_uid):
        """Invoke processing of the window."""
        output_msgs = self._processor.process(
            self._window_state_map[window_uid])
        for output_msg in output_msgs:
            self.get_output_stream(self._output_stream_name).send(output_msg)
        # Remove window state
        del self._window_state_map[window_uid]
        del self._window_trigger_map[window_uid]
        try:
            self._registered_windows.remove(window_uid)
        except KeyError as e:
            pass

    def on_time_trigger(self, time, window):
        # The time window ended, invoke trigger to decide action.
        action = self._trigger.on_time(time, window)
        if action == WindowTriggerAction.PROCESS:
            self.on_process(window.uid)

    def on_message_trigger(self, msg, window):
        action = self._window_trigger_map[window.uid].on_message(msg, window)
        if action == WindowTriggerAction.PROCESS:
            self.on_process(window.uid)

    def register_time_trigger(self, window_end_time, window):
        """Adds the end of the window to a priority queue."""
        if window.uid not in self._registered_windows:
            # Only add the end of window event if we haven't seen this window
            # before.
            self._registered_windows.add(window.uid)
            heapq.heappush(self._window_end_pqueue, (window_end_time, window))

    def on_msg(self, msg):
        windows = self._assigner.assign_windows(msg)
        for window in windows:
            # Add message to the per-window state.
            if window.uid in self._window_state_map:
                self._window_state_map[window.uid].append(msg)
            else:
                self._window_state_map[window.uid] = [msg]
            if window.uid not in self._window_trigger_map:
                self._window_trigger_map[window.uid] = copy.deepcopy(
                    self._trigger)
            # TODO(ionel): We shouldn't have two paths for time & non-time
            # windows. Fix!
            if window.is_time_window:
                # Register a trigger for the end of the window.
                self.register_time_trigger(window.end_time, window)
            self.on_message_trigger(msg, window)

    @frequency(100)
    def fire_triggers(self):
        # TODO(ionel): There are more efficient ways to implement this. Fix!
        while (len(self._window_end_pqueue) > 0
               and self._window_end_pqueue[0][0] <= get_timestamp_ms()):
            (end_time, window) = heapq.heappop(self._window_end_pqueue)
            self.on_time_trigger(end_time, window)

    def execute(self):
        self.fire_triggers()
        self.spin()


class NoopOp(Op):
    def on_msg(self, msg):
        self.get_output_stream(msg.stream_name).send(msg)

    @staticmethod
    def setup_streams(input_streams, **kwargs):
        """Note: does not support duplicate stream names"""
        stream_names = set()
        output_streams = []
        for input_stream in input_streams._streams:
            # Make sure no 2 streams have identical names
            assert input_stream.name not in stream_names, "NoopOp does not support multiple input streams with identical names"
            stream_names.add(input_stream.name)
            output_stream = DataStream(input_stream.data_type,
                                       input_stream.name, input_stream.labels)
            output_streams.append(output_stream)

        input_streams.add_callback(NoopOp.on_msg)

        return output_streams

# We need the release operators with the following behaviors:
# - Releases the timestamp message at the same time and releases the
#   remaining messages from the same timestamp from the same stream.
# - Waits until the deadline is over and then releases the results.


class EarlyReleaseOperator(Op):
    def __init__(self, name, output_stream_name):
        super(EarlyReleaseOperator, self).__init__(name)
        self.output_stream_name = output_stream_name
        self.timestamp_to_stream_name_map = {}

    def on_msg(self, msg):
        stream_name = self.timestamp_to_stream_name_map.get(msg.timestamp)
        if stream_name:
            # Stream name exists, so we have seen a message from the same
            # timestamp before. Release if it is from the same stream.
            if stream_name == msg.stream_name:
                self.get_output_stream(self.output_stream_name).send(msg)
        else:
            # Stream name does not exist. Add it to the map and release this
            # message
            self.timestamp_to_stream_name_map[msg.timestamp] = msg.stream_name
            self.get_output_stream(self.output_stream_name).send(msg)

    def on_completion_msg(self, msg):
        stream_name = self.timestamp_to_stream_name_map.get(msg.timestamp)
        if stream_name:
            # The stream name exists. We should check if the watermark is from
            # the stream. If yes, release, otherwise let it be.
            if stream_name == msg.stream_name:
                # TODO (sukritk) FIX (Ray #4463): Remove when Ray issue fixed.
                watermark_msg = WatermarkMessage(msg.timestamp,
                                                 msg.stream_name)
                self.get_output_stream(
                    self.output_stream_name).send(watermark_msg)

    @staticmethod
    def setup_streams(input_streams, output_stream_name):
        if len(input_streams) > 0:
            assert all(input_stream.data_type == input_streams[0].data_type \
                       for input_stream in input_streams), \
                "All input streams into the ReleaseOperator " \
                "should have the same data type"
            input_streams.add_callback(EarlyReleaseOperator.on_msg)
            input_streams.add_completion_callback(
                EarlyReleaseOperator.on_completion_msg)
            return [
                DataStream(input_streams[0].data_type, output_stream_name,
                           input_streams[0].labels)
            ]
        else:
            return []
