try: from _thread import start_new_thread
except: from thread import start_new_thread

import os, sys
import logging
import traceback
import time
from multiprocessing import Process

from jobqueue import runjob, store
from jobqueue.client import connect

store.ROOT = '/tmp/worker/%s'
DEFAULT_DISPATCHER = 'http://reflectometry.org/queue'
POLLRATE = 10

def log_errors(f):
    def wrapped(*args, **kw):
        try:
            return f(*args, **kw)
        except:
            exc_type,exc_value,exc_trace = sys.exc_info()
            trace = traceback.format_tb(exc_trace)
            message = traceback.format_exception_only(exc_type,exc_value)
            logging.error(message+trace)
    return wrapped

def wait_for_result(remote, id, process, queue):
    """
    Wait for job processing to finish.  Meanwhile, prefetch the next
    request.
    """
    next_request = { 'request': None }
    canceling = False
    while True:
        # Check if process is complete
        process.join(POLLRATE)
        if not process.is_alive(): break

        # Check that the job is still active, and that it hasn't been
        # canceled, or results reported back from a second worker.
        # If remote server is down, assume the job is still active.
        try: response = remote.status(id)
        except: response = None
        if response and response['status'] != 'ACTIVE':
            #print "canceling process"
            process.terminate()
            canceling = True
            break

        # Prefetch the next job; this strategy works well if there is
        # only one worker.  If there are many, we may want to leave it
        # for another worker to process.
        if not next_request['request']:
            # Ignore remote server down errors
            try: next_request = remote.nextjob(queue=queue)
            except: pass

    # Grab results from the store
    try:
        results = runjob.results(id)
    except KeyError:
        if canceling:
            results = { 'status': 'CANCEL', 'message': 'Job canceled' }
        else:
            results = { 'status': 'ERROR', 'message': 'Results not found' }

    #print "returning results",results
    return results, next_request

@log_errors
def update_remote(dispatcher, id, queue, results):
    """
    Update remote server with results.
    """
    #print "updating remote"
    path= store.path(id)
    # Remove results key, if it is there
    try: store.delete(id, 'results')
    except KeyError: pass
    files = [os.path.join(path,f) for f in os.listdir(path)]
    #print "sending results",results
    # This is done with a separate connection to the server so that it can
    # run inside a thread.  That way the server can start the next job
    # while the megabytes of results are being transfered in the background.
    private_remote = connect(dispatcher)
    private_remote.postjob(id=id, results=results, queue=queue, files=files)
    # Clean up files
    for f in files: os.unlink(f)
    os.rmdir(path)

def serve(dispatcher, queue):
    """
    Run the work server.
    """
    assert queue is not None
    next_request = { 'request': None }
    remote = connect(dispatcher)
    while True:
        if not next_request['request']:
            try: next_request = remote.nextjob(queue=queue)
            except: logging.error(traceback.format_exc())
        if next_request['request']:
            jobid = next_request['id']
            if jobid is None:
                logging.error('request has no job id')
                next_request = {'request': None}
                continue
            logging.info('processing job %s'%jobid)
            process = Process(target=runjob.run,
                              args=(jobid,next_request['request']))
            process.start()
            results, next_request = wait_for_result(remote, jobid, process, queue)
            start_new_thread(update_remote,
                             (dispatcher, jobid, queue, results))
        else:
            time.sleep(POLLRATE)

def main():
    try: os.nice(19)
    except: pass
    if len(sys.argv) <= 1:
        print("Requires queue name")
    queue = sys.argv[1]
    dispatcher = sys.argv[2] if len(sys.argv) > 2 else DEFAULT_DISPATCHER
    serve(queue=queue, dispatcher=dispatcher)

if __name__ == "__main__":
    main()