#  -*- coding: utf-8 -*-
# *****************************************************************************
# ufit, a universal scattering fitting suite
#
# Copyright (c) 2013-2023, Georg Brandl and contributors.  All rights reserved.
# Licensed under a 2-clause BSD license, see LICENSE.
# *****************************************************************************

"""
Utilities for clustering execution of a piece of code to multiple hosts using
SSH transport.
"""

import hashlib
import pickle
import queue
import sys
import threading
from os import path
from time import time

import paramiko

keyname = path.expanduser('~/.ufitcluster/key')
clusterlist = []


def init_cluster():
    try:
        with open(path.expanduser('~/.ufitcluster/hosts'),
                  encoding='utf-8') as fp:
            for line in fp:  # user@host
                if line.startswith('#'):
                    continue
                login, host = line.strip().split('@')
                clusterlist.append((login, host))
    except Exception as err:
        print('Cannot read the list of cluster hosts:', err, file=sys.stderr)
        print('Please create a file %s with one line for each '
              'parallel process, for example:' %
              path.expanduser('~/.ufitcluster/hosts'), file=sys.stderr)
        print(file=sys.stderr)
        print('username@hostname1', file=sys.stderr)
        print('username@hostname1', file=sys.stderr)
        print('username@hostname1', file=sys.stderr)
        print('username@hostname1', file=sys.stderr)
        print('username@hostname2', file=sys.stderr)
        print('username@hostname2', file=sys.stderr)
        print(file=sys.stderr)
        print('would run 4 parallel processes on hostname1 and '
              '2 parallel processes on hostname2.', file=sys.stderr)
        print('For local operation only, use "localhost" as '
              'hostname.', file=sys.stderr)
        print('The file %s must contain the SSH private key to '
              'use for the connection, even for localhost.' %
              path.expanduser('~/.ufitcluster/key'), file=sys.stderr)


def client_runner(client, task_queue, result_queue):
    while True:
        code, funcname, jobnum, args = task_queue.get()
        if jobnum == -1:
            print('[C] exiting runner for %s' % client.ufit_host)
            return
        try:
            sidcontents = (str(time()) + str(args)).encode()
            sid = hashlib.md5(sidcontents).hexdigest()
            sys.stdout.write('.')
            sys.stdout.flush()
            # print '[C] starting job on %s: %s' % (client.ufit_host, sid)
            code_footer = '''\nif __name__ == "__main__":
            import pickle
            args = pickle.loads(%r)
            print pickle.dumps(%s(*args))
            ''' % (pickle.dumps(args), funcname)

            sftp = client.open_sftp()
            fobj = sftp.file('/tmp/ufit_cluster_%s.py' % sid, 'wb')
            fobj.write(code + code_footer)
            fobj.close()

            _stdin, stdout, _stderr = \
                client.exec_command('python /tmp/ufit_cluster_%s.py; '
                                    'rm /tmp/ufit_cluster_%s.py' % (sid, sid))
            result = pickle.load(stdout)
            # print '[C] done with job %s: %r' % (sid, result)
            result_queue.put((jobnum, result))
        except Exception as err:
            print(f'[C] no result on {client.ufit_host}, requeuing: {err!r}')
            task_queue.put((code, funcname, jobnum, args))
            return


clients = []
runners = []
global_task_queue = queue.Queue()
global_result_queue = queue.Queue()
cluster_setup = False


def setup_cluster():
    global cluster_setup
    if cluster_setup:
        return
    for i, (user, host) in enumerate(clusterlist):
        client = paramiko.SSHClient()
        client.ufit_host = '%s[%d]' % (host, i)
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        client.connect(host, username=user, key_filename=keyname,
                       look_for_keys=False)
        clients.append(client)
        print(f'[C] opened cluster connection to {client.ufit_host}')
        runner = threading.Thread(target=client_runner,
                                  args=(client, global_task_queue,
                                        global_result_queue))
        runners.append(runner)
        runner.start()
    cluster_setup = True


def kill_cluster():
    for cl in clients:
        global_task_queue.put((None, None, -1, None))  # end!
        cl.close()


def run_cluster(code, funcname, argumentslist):
    setup_cluster()
    njobs = len(argumentslist)
    retval = [None] * njobs
    returns = 0
    for job in enumerate(argumentslist):
        global_task_queue.put((code, funcname) + job)
    while returns < njobs:
        # XXX check if any client is still running
        jobnum, result = global_result_queue.get()
        if isinstance(result, Exception):
            raise result
        retval[jobnum] = result
        returns += 1
    return retval


init_cluster()

if __name__ == '__main__':
    t1 = time()
    print(run_cluster('def foo(a): return a\n', 'foo', [('a',)]))
    t2 = time()
    print(t2-t1)
