import datetime

import itertools
import random

import zmq
import time

from motorway.exceptions import SocketBlockedException
from motorway.grouping import HashRingGrouper, RandomGrouper, GroupingValueMissing, SendToAllGrouper
from motorway.messages import Message
from motorway.utils import set_timeouts_on_socket, get_connections_block, get_ip


class GrouperMixin(object):
    def get_grouper(self, grouper_name):
        if grouper_name == "HashRingGrouper":
            return HashRingGrouper
        elif grouper_name == "SendToAllGrouper":
            return SendToAllGrouper
        else:
            return RandomGrouper


class SendMessageMixin(object):

    def send_message(self, message, process_id, time_consumed=None, sender=None, control_message=True, retries=100):
        """

        :param message: A ::pipeline.messages.Message:: instance
        :param process_id:
        :param time_consumed:
        :param sender:
        :param control_message:
        :param retries:
        :type retries: int
        :raises: SocketBlockedException, GroupingValueMissing
        :return:
        """
        try:
            socket_addresses = self.grouper_instance.get_destinations_for(message.grouping_value)
        except GroupingValueMissing:
            raise GroupingValueMissing("Message '%s' provided an invalid grouping_value: '%s'" % (message.content, message.grouping_value))
        for destination in socket_addresses:
            message.send(
                self.send_socks[destination],
                process_id
            )
            if self.controller_sock and self.send_control_messages and control_message:
                for index in range(0, retries):
                    try:
                        message.send_control_message(
                            self.controller_sock,
                            time_consumed=time_consumed,
                            process_name=process_id,
                            destination_uuid=self.process_address_to_uuid[destination],
                            destination_endpoint=destination,
                            sender=sender
                        )
                        break
                    except KeyError:  # If destination is not known, lookup fails
                        time.sleep(random.randrange(1, 3))  # random to avoid peak loads when running multiple processes
                    except zmq.Again:  # If connection fails, retry
                        time.sleep(random.randrange(1, 3))
                else:  # if for loop exited cleanly (no break)
                    raise SocketBlockedException("Control Message for process %s could not be sent after %s attempts"
                                                 % (process_id, retries))


class ConnectionMixin(object):
    def set_send_socks(self, connections, output_queue, context):
        """
        Sets the sockets for which any outgoing messages should be sent to.

        :param connections: dict, containing all available queues and streams
        :param output_queue: str, name of the queue to output messages to
        :param context: 'zmq.Context'-instance
        """
        for send_conn in connections[output_queue]['streams']:
            if send_conn not in self.send_socks:
                send_sock = context.socket(zmq.PUSH)
                send_sock.connect(send_conn)
                self.send_socks[send_conn] = send_sock
                self.send_grouper = connections[output_queue]['grouping']

        deleted_connections = [connection for connection in self.send_socks.keys() if connection not in connections[output_queue]['streams']]
        for deleted_connection in deleted_connections:
            self.send_socks[deleted_connection].close()
            del self.send_socks[deleted_connection]

        # initialize grouper again with new socks
        self.grouper_instance = self.get_grouper(self.send_grouper)(
            self.send_socks.keys()
        )

    def connection_thread(self, context=None, refresh_connection_stream=None, input_queue=None, output_queue=None,
                          grouper_cls=None, set_controller_sock=True):
        """
        Responsible for managing and connecting to 'input_queue' and 'output_queue' for a ramp or intersection.

        It is also responsible for sending heartbeats to the ConnectionIntersection and marking itself
        as a consumer of an 'input_queue', which can be either a ramp or intersection.

        :param context: 'zmq.Context'-instance
        :param refresh_connection_stream: str, TCP-address of the ConnectionIntersection
        :param input_queue: str, name of the queue to receive messages from
        :param output_queue: str, name of the queue to output messages to
        :param grouper_cls: (Optional) Grouper-class inheriting from BaseGrouper
        :param set_controller_sock: bool, determines whether or not to track messages in the ControllerIntersection
        """
        refresh_connection_sock = context.socket(zmq.SUB)
        set_timeouts_on_socket(refresh_connection_sock)
        refresh_connection_sock.connect(refresh_connection_stream)
        refresh_connection_sock.setsockopt_string(zmq.SUBSCRIBE, u'')  # You must subscribe to something, so this means *all*}

        connections = get_connections_block('_update_connections', refresh_connection_sock)

        while not self.receive_port:
            time.sleep(1)

        # Register as consumer of input stream
        update_connection_sock = context.socket(zmq.PUSH)
        set_timeouts_on_socket(update_connection_sock)
        update_connection_sock.connect(connections['_update_connections']['streams'][0])
        intersection_connection_info = {
            'streams': {
                input_queue: ['tcp://%s:%s' % (get_ip(), self.receive_port)]
            },
            'meta': {
                'id': self.process_uuid,
                'name': self.process_name,
                'grouping': None if not grouper_cls else grouper_cls.__name__  # Specify how messages sent to this intersection should be grouped
            }
        }
        Message(None, intersection_connection_info).send(update_connection_sock, producer_uuid=self.process_uuid)
        last_send_time = datetime.datetime.now()

        if set_controller_sock:
            connections = get_connections_block('_message_ack', refresh_connection_sock, existing_connections=connections)
            self.controller_sock = context.socket(zmq.PUSH)
            set_timeouts_on_socket(self.controller_sock)
            self.controller_sock.connect(connections['_message_ack']['streams'][0])
        while True:
            try:
                connections = refresh_connection_sock.recv_json()
                if output_queue in connections:
                    self.set_send_socks(connections, output_queue, context)
                self.process_id_to_name = {process['process_id']: process['process_name'] for process in itertools.chain.from_iterable([queue_details['stream_heartbeats'].values() for queue_details in connections.values()])}
                self.process_address_to_uuid = {address: process['process_id'] for address, process in itertools.chain.from_iterable([queue_details['stream_heartbeats'].items() for queue_details in connections.values()])}
            except zmq.Again:
                time.sleep(1)
            if last_send_time + datetime.timedelta(seconds=10) < datetime.datetime.now():
                Message(None, intersection_connection_info).send(update_connection_sock, producer_uuid=self.process_uuid)
                last_send_time = datetime.datetime.now()
