from __future__ import print_function

import os
import itertools
import uuid
from datetime import datetime

import zmq
import dill
from .scheduler import pickle
from ..compatibility import unicode

context = zmq.Context()

jobids = ('schedule-%d' % i for i in itertools.count())


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

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



class Client(object):
    def __init__(self, scheduler, address=None):
        self.address_to_scheduler = scheduler
        if address == None:
            address = 'client-' + str(uuid.uuid1())
        if isinstance(address, unicode):
            address = address.encode()
        self.address = address
        self.socket = context.socket(zmq.DEALER)
        self.socket.setsockopt(zmq.IDENTITY, self.address)
        self.socket.connect(self.address_to_scheduler)
        self.register_client()

    def get(self, dsk, keys):
        header = {'function': 'schedule',
                  'jobid': next(jobids)}
        payload = {'dask': dsk, 'keys': keys}

        self.send_to_scheduler(header, payload)
        header2, payload2 = self.recv_from_scheduler()

        if header2['status'] != 'OK':
            raise payload2['result']

        return payload2['result']

    def scheduler_status(self):
        header = {'function': 'status'}
        payload = {}
        self.send_to_scheduler(header, payload)

        header2, payload2 = self.recv_from_scheduler()
        return payload2

    def send_to_scheduler(self, header, payload):
        log(self.address, 'Send to scheduler', header)
        if 'address' not in header:
            header['address'] = self.address
        header['timestamp'] = datetime.utcnow()
        header['loads'] = dill.loads
        self.socket.send_multipart([pickle.dumps(header), dill.dumps(payload)])

    def recv_from_scheduler(self):
        header, payload = self.socket.recv_multipart()
        header = pickle.loads(header)
        loads = header.get('loads', pickle.loads)
        payload = loads(payload)
        log(self.address, 'Received from scheduler', header)
        return header, payload

    def set_collection(self, name, collection):
        """ Store collection in scheduler

        See docstring for get_collection
        """
        header = {'function': 'set-collection',
                  'loads': dill.loads}
        payload = {'type': type(collection),
                   'args': collection._args,
                   'name': name}
        self.send_to_scheduler(header, payload)
        header2, payload2 = self.recv_from_scheduler()

        assert header2['status'] == 'OK'

    def get_collection(self, name):
        """ Get stored collection from scheduler

        Clients may share collections with other clients by registering them
        with the centralized scheduler.

        >>> import dask.bag as db  # doctest: +SKIP
        >>> b = db.from_sequence(...).map(...).filter(...) # doctest: +SKIP

        >>> from dask.distributed import Client  # doctest: +SKIP
        >>> client = Client('tcp://scheduler-hostname:5555')  # doctest: +SKIP
        >>> client.set_collection('mybag', b)  # doctest: +SKIP

        Other clients can connect to the same scheduler to collect that
        collection.

        >>> client2 = Client('tcp://scheduler-hostname:5555')  # doctest: +SKIP
        >>> b2 = client2.get_collection('mybag')  # doctest: +SKIP
        """
        header = {'function': 'get-collection'}
        payload = {'name': name}

        self.send_to_scheduler(header, payload)
        header2, payload2 = self.recv_from_scheduler()

        return payload2['type'](*payload2['args'])

    def close_scheduler(self):
        header = {'function': 'close'}
        self.send_to_scheduler(header, {})

    def close(self, close_scheduler=False):
        if close_scheduler:
            self.close_scheduler()
        self.socket.close(1)

    def register_client(self):
        self.send_to_scheduler({'function': 'register'}, {'pid': os.getpid()})
        header, payload = self.recv_from_scheduler()
        self.registered_workers = payload['workers']

    def get_registered_workers(self):
        self.send_to_scheduler({'function': 'get_workers'}, {})
        header, payload = self.recv_from_scheduler()
        return payload['workers']

@@ -2178,16 +2178,149 @@
 ayload%0A%0A
+    def send_recv(self, header, payload):%0A        self.send_to_scheduler(header, payload)%0A        return self.recv_from_scheduler()%0A%0A
     def 
@@ -2359,16 +2359,16 @@
 ction):%0A
-
         
@@ -2669,55 +2669,8 @@
 me%7D%0A
-        self.send_to_scheduler(header, payload)
 %0A   
@@ -2699,36 +2699,41 @@
  = self.
-recv_from_scheduler(
+send_recv(header, payload
 )%0A%0A     
@@ -3650,56 +3650,8 @@
 e%7D%0A%0A
-        self.send_to_scheduler(header, payload)%0A
     
@@ -3679,36 +3679,41 @@
  = self.
-recv_from_scheduler(
+send_recv(header, payload
 )%0A%0A     
@@ -4049,39 +4049,25 @@
         
-self.send_to_scheduler(
+header = 
 %7B'functi
@@ -4082,17 +4082,34 @@
 gister'%7D
-,
+%0A        payload =
  %7B'pid':
@@ -4121,17 +4121,16 @@
 etpid()%7D
-)
 %0A       
@@ -4128,32 +4128,33 @@
 %7D%0A        header
+2
 , payload = self
@@ -4142,32 +4142,33 @@
 er2, payload
+2
  = self.
 recv_from_sc
@@ -4155,36 +4155,41 @@
  = self.
-recv_from_scheduler(
+send_recv(header, payload
 )%0A      
@@ -4215,32 +4215,33 @@
 orkers = payload
+2
 %5B'workers'%5D%0A%0A   

Add Client.send_recv method.