# Darwin Workers behave like servers.
# They respond to assignation request from the main node.

import logging
import sys
import zmq
from utils import (
    darwin_send_data,
    darwin_receive_data,
    darwin_send_signal,
    darwin_receive_signal
)

logging.basicConfig(
    format='[%(levelname)s][%(asctime)s]: %(message)s', level=logging.INFO)


class DarwinWorker(object):
    """Darwin default Worker (should subclass and redefine `work` method"""

    manifest = {
        'name': 'Default Worker',
        'code': 'DEFAULT',
        'port': '7770'
    }

    def __init__(self, manifest={}):
        """Set some init variables."""
        # ZeroMQ variables
        self.context = None
        self.socket = None

        # Assignation variables
        self.core_data = None
        self.assignation = None
        self.result = None

        # Update default manifest with passed one
        self.manifest.update(manifest)

        # Shortcuts
        self.name = self.manifest['name']
        self.port = self.manifest['port']
        self.code = self.manifest['code']

    def wait_for_work(self):
        """Initialize the worker and waits for any request."""

        # Setup ZeroMQ and other things
        self.initialize_connection()

        logging.info("Waiting for work...")
        while True:
            # Wait for signal (Blocks on recv)
            signal = darwin_receive_signal(self.socket)

            # If START signal... WORK!
            if signal == 'START':
                self.start()
            else:
                logging.warning("Unknown signal: [ %s ]" % signal)

    def initialize_connection(self):
        """Sets up ZeroMQ and a socket for this worker."""

        logging.info("Initializing ZeroMQ...")
        self.context = zmq.Context()

        logging.info("Initializing socket...")
        self.socket = self.context.socket(zmq.REP)

        logging.info("Starting Algorithm: [ %s ] at port %s" % (
            self.name, self.port))
        self.socket.bind("tcp://*:%s" % self.port)

    def start(self):
        """Starts the execution cycle for the worker."""

        # Started. Send an 'OK'
        darwin_send_signal(self.socket, 'OK')

        # Get data
        logging.info("Receiving Core Data...")
        self.core_data = darwin_receive_data(self.socket)

        # Send an 'OK' an receive the assignation data
        darwin_send_signal(self.socket, 'OK')
        logging.info("Receiving Assignation Data...")
        self.assignation = darwin_receive_data(self.socket)

        # Send a 'READY' an start working...
        logging.info("Ready to work")
        darwin_send_signal(self.socket, 'READY')

        # Esperamos un OK del Master, y empezamos a trabajar...
        logging.info("Getting confirmation from Master...")
        response_signal = darwin_receive_signal(self.socket)

        if response_signal == 'OK':
            logging.info('WORK!!')

            # Let's invoke worker method and fetch its result
            self.result = self.work(self.core_data, self.assignation)

        # Now, let's end this...
        self.end()

    def end(self):
        # Signal the master that we are ready...
        logging.info('Done Working. Ending negotiation...')
        darwin_send_signal(self.socket, 'DONE')

        # Get an 'OK'
        signal_response = darwin_receive_signal(self.socket)

        if signal_response == 'OK':
            # It's ok to send data...
            darwin_send_data(self.socket, self.result)

            # Wait for reply...
            signal_response = darwin_receive_signal(self.socket)

            if signal_response == 'OK':
                # Master got the data
                # so let's signal that we finished working....
                darwin_send_signal(self.socket, 'FINISHED')
                logging.info('Done Working d^_^b')

    def work(self, base, assignation):
        """Work method for the Algorithm. You should override this method.
        You must return a new assignation
        in the same format as the one received"""
        return assignation


if __name__ == '__main__':
    if len(sys.argv) > 1:
        port = sys.argv[1]
        int(port)

    w = DarwinWorker()

    w.wait_for_work()
