from __future__ import print_function

import socket
from threading import Thread, Lock
from multiprocessing.pool import ThreadPool
from contextlib import contextmanager
import traceback
from datetime import datetime
import uuid
import random
import multiprocessing
import zmq
from time import time
import sys
from ..compatibility import Queue, unicode
from .. import core
try:
    import cPickle as pickle
except ImportError:
    import pickle
import dill

def pickle_dumps(obj):
    return pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)

MAX_DEALERS = 100

with open('log.workers', 'w') as f:  # delete file
    pass

def log(*args):
    with open('log.workers', 'a') as f:
        print(*args, file=f)

log('Hello from worker.py')

@contextmanager
def logerrors():
    try:
        yield
    except Exception as e:
        exc_type, exc_value, exc_traceback = sys.exc_info()
        tb = ''.join(traceback.format_tb(exc_traceback))
        log('Error!', str(e))
        log('Traceback', str(tb))
        raise

class Worker(object):
    """ Asynchronous worker in a distributed dask computation pool


    Parameters
    ----------

    scheduler: string
        Address of scheduler
    hostname: string
        A visible hostname/IP of this worker to the network
    port_to_workers: int
        Port on which to listen to worker connections
    bind_to_workers: string
        Addresses from which we accept worker connections, defaults to *

    State
    -----

    status: string
        Status of worker, either 'run' or 'closed'
    to_workers: zmq.Socket
        Router socket to serve requests from other workers
    to_scheduler: zmq.Socket
        Dealer socket to communicate with scheduler

    See Also
    --------

        dask.distributed.scheduler.Scheduler
    """
    def __init__(self, scheduler, data=None, nthreads=100,
                 hostname=None, port_to_workers=None, bind_to_workers='*',
                 block=False):
        if isinstance(scheduler, unicode):
            scheduler = scheduler.encode()
        self.data = data if data is not None else dict()
        self.pool = ThreadPool(nthreads)
        self.scheduler = scheduler
        self.status = 'run'
        self.context = zmq.Context()

        self.hostname = hostname or socket.gethostname()

        self.to_workers = self.context.socket(zmq.ROUTER)
        if port_to_workers is None:
            port_to_workers = self.to_workers.bind_to_random_port('tcp://' + bind_to_workers)
        else:
            self.to_workers.bind('tcp://%s:%d' % (bind_to_workers, port))
        self.address = ('tcp://%s:%s' % (self.hostname, port_to_workers)).encode()

        self.dealers = dict()

        self.lock = Lock()

        self.queues = dict()

        self.to_scheduler = self.context.socket(zmq.DEALER)

        self.to_scheduler.setsockopt(zmq.IDENTITY, self.address)
        self.to_scheduler.connect(scheduler)
        self.send_to_scheduler({'function': 'register'}, {})

        self.scheduler_functions = {'status': self.status_to_scheduler,
                                    'compute': self.compute,
                                    'getitem': self.getitem_scheduler,
                                    'delitem': self.delitem,
                                    'setitem': self.setitem,
                                    'close': self.close_from_scheduler}

        self.worker_functions = {'getitem': self.getitem_worker,
                                 'getitem-ack': self.getitem_ack,
                                 'status': self.status_to_worker}

        log(self.address, 'Start up', self.scheduler)

        self._listen_scheduler_thread = Thread(target=self.listen_to_scheduler)
        self._listen_scheduler_thread.start()
        self._listen_workers_thread = Thread(target=self.listen_to_workers)
        self._listen_workers_thread.start()

        if block:
            self.block()

    def status_to_scheduler(self, header, payload):
        out_header = {'jobid': header.get('jobid')}
        log(self.address, 'Status check', header['address'])
        self.send_to_scheduler(out_header, 'OK')

    def status_to_worker(self, header, payload):
        out_header = {'jobid': header.get('jobid')}
        log(self.address, 'Status check', header['address'])
        self.send_to_worker(header['address'], out_header, 'OK')

    def getitem_worker(self, header, payload):
        """ Get data and send to another worker

        See also:
            Worker.collect
        """
        loads = header.get('loads', pickle.loads)
        payload = loads(payload)
        log(self.address, "Getitem for worker", header, payload)
        header2 = {'function': 'getitem-ack',
                   'jobid': header.get('jobid')}
        try:
            result = self.data[payload['key']]
            header2['status'] = 'OK'
        except KeyError as e:
            result = e
            header2['status'] = 'Bad key'
        payload = {'key': payload['key'],
                   'value': result,
                   'queue': payload['queue']}
        self.send_to_worker(header['address'], header2, payload)

    def getitem_ack(self, header, payload):
        """ Receive data after sending a getitem request

        See also:
            Worker.getitem_worker
            Worker.collect
        """
        with logerrors():
            loads = header.get('loads', pickle.loads)
            payload = loads(payload)
            log(self.address, 'Getitem ack', payload)
            assert header['status'] == 'OK'

            self.data[payload['key']] = payload['value']
            self.queues[payload['queue']].put(payload['key'])

    def getitem_scheduler(self, header, payload):
        """ Send local data to scheduler

        See also:
            Scheduler.gather
            Scheduler.getitem_ack
        """
        loads = header.get('loads', pickle.loads)
        payload = loads(payload)
        log(self.address, 'Get from scheduler', payload)
        key = payload['key']
        header2 = {'jobid': header.get('jobid')}
        try:
            result = self.data[key]
            header2['status'] = 'OK'
        except KeyError as e:
            result = e
            header2['status'] = 'Bad key'
        header2['function'] = 'getitem-ack'
        payload2 = {'key': key, 'value': result, 'queue': payload['queue']}
        self.send_to_scheduler(header2, payload2)

    def setitem(self, header, payload):
        """ Assign incoming data to local dictionary

        See also:
            Scheduler.scatter
            Scheduler.send_data
            Scheduler.setitem_ack
        """
        loads = header.get('loads', pickle.loads)
        payload = loads(payload)
        log(self.address, 'Setitem', payload)
        key = payload['key']
        value = payload['value']
        self.data[key] = value

        queue = payload.get('queue', False)
        if queue:
            header2 = {'jobid': header.get('jobid'),
                       'function': 'setitem-ack'}
            payload2 = {'key': key, 'queue': queue}
            log(self.address, 'Setitem send ack to scheduler',
                header2, payload2)
            self.send_to_scheduler(header2, payload2)

    def delitem(self, header, payload):
        """ Remove item from local data """
        loads = header.get('loads', pickle.loads)
        payload = loads(payload)
        log(self.address, 'Delitem', payload)
        key = payload['key']
        del self.data[key]

        # TODO: this should be replaced with a delitem-ack call
        if payload.get('reply', False):
            self.send_to_scheduler({'jobid': header.get('jobid')}, 'OK')


    def send_to_scheduler(self, header, payload):
        """ Send data to scheduler """
        log(self.address, 'Send to scheduler', header)
        header['address'] = self.address
        header['timestamp'] = datetime.utcnow()
        dumps = header.get('dumps', pickle_dumps)
        with self.lock:
            self.to_scheduler.send_multipart([pickle_dumps(header),
                                              dumps(payload)])

    def send_to_worker(self, address, header, payload):
        """ Send data to workers

        This is a bit tricky.  We want to have one DEALER socket per worker.
        We cache these in ``self.dealers``.  If the number of worker peers is
        high then we might run into having too many file descriptors open.
        Currently we flush the cache of dealers periodically.  This has yet to
        be tested.
        """
        if address not in self.dealers:
            if len(self.dealers) > MAX_DEALERS:
                for sock in self.dealers.values():
                    sock.close()
                self.dealers.clear()
            sock = self.context.socket(zmq.DEALER)
            sock.connect(address)
            self.dealers[address] = sock

        header['address'] = self.address
        header['timestamp'] = datetime.utcnow()
        log(self.address, 'Send to worker', address, header)
        dumps = header.get('dumps', pickle_dumps)
        with self.lock:
            self.dealers[address].send_multipart([pickle_dumps(header),
                                                  dumps(payload)])

    def listen_to_scheduler(self):
        """
        Event loop listening to commands from scheduler

        Header and Payload should deserialize into dicts of the following form:

            Header
            {'function': name of function to call, see self.functions,
             'jobid': job identifier, defaults to None,
             'address': name of sender, defaults to zmq identity}

            Payload
            --Function specific, for setitem might include the following--
            {'key': 'x',
             'value': 10}

        So the minimal request would be as follows:

        >>> sock = context.socket(zmq.DEALER)  # doctest: +SKIP
        >>> sock.connect('tcp://my-address')   # doctest: +SKIP

        >>> header = {'function': 'status'}
        >>> payload = {}
        >>> sock.send_multipart(dumps(header), dumps(status))  # doctest: +SKIP

        Or a more complex packet might be as follows:

        >>> header = {'function': 'setitem', 'jobid': 1}
        >>> payload = {'key': 'x', 'value': 10}
        >>> sock.send_multipart(dumps(header), dumps(status))  # doctest: +SKIP

        We match the function string against ``self.scheduler_functions`` to
        pull out the actual function.  We then execute this function with the
        provided arguments in another thread from ``self.pool``.  That function
        may then choose to send results back to the sender.

        See Also:
            listen_to_workers
            send_to_scheduler
        """
        while self.status != 'closed':
            # Wait on request
            try:
                if not self.to_scheduler.poll(100):
                    continue
            except zmq.ZMQError:
                break
            with logerrors():
                with self.lock:
                    header, payload = self.to_scheduler.recv_multipart()
                header = pickle.loads(header)
                log(self.address, 'Receive job from scheduler', header)