import datetime

import itertools
import random
import socket
import uuid

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.get_grouper(self.send_grouper)(
                self.send_socks.keys()
            ).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 xrange(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):
        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]

    def connection_thread(self, context=None, refresh_connection_stream=None, input_queue=None, output_queue=None,
                          grouper_cls=None, set_controller_sock=True):
        refresh_connection_sock = context.socket(zmq.SUB)
        refresh_connection_sock.connect(refresh_connection_stream)
        refresh_connection_sock.setsockopt(zmq.SUBSCRIBE, '')  # You must subscribe to something, so this means *all*}
        set_timeouts_on_socket(refresh_connection_sock)

        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)
        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)
            self.controller_sock.connect(connections['_message_ack']['streams'][0])
            set_timeouts_on_socket(self.controller_sock)
        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()
