#!/usr/bin/env python

import logging
from sys import argv, exit
import socket

import click
from distributed import Nanny, Worker, sync
from distributed.utils import get_ip
from tornado.ioloop import IOLoop

logger = logging.getLogger('distributed.dworker')

ip = get_ip()


@click.command()
@click.argument('center', type=str)
@click.option('--port', type=int, default=8788, help="Serving port")
@click.option('--host', type=str, default=ip, help="Serving host")
@click.option('--ncores', type=int, default=0, help="Number of cores to serve")
def go(center, host, port, ncores):
    while True:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            sock.bind(('127.0.0.1', port))
        except socket.error:
            port += 1
        else:
            sock.close()
            break

    try:
        center_ip, center_port = center.split(':')
        center_port = int(center_port)
    except IndexError:
        logger.info("Usage:  dworker center_host:center_port")

    loop = IOLoop.current()
    worker = Nanny(host, port, port + 1, center_ip, center_port, ncores=ncores)
    loop.add_callback(worker._start)
    try:
        loop.start()
    except KeyboardInterrupt:
        result = IOLoop().run_sync(lambda:
                worker.center.unregister(address=worker.worker_address))
        if result == b'OK':
            logger.info("Unregister %s:%d from center", *worker.worker_address)
        logger.info("End worker at %s:%d", ip, port)

if __name__ == '__main__':
    go()
