#!/usr/bin/env python2.6
import sys
import logging
import threading
import Queue
from threadpool import ThreadPool
from task import Task, ThreadedTask

class Scheduler(threading.Thread):

    class _SchedulerStop(Exception):
        pass

    def __init__(self, poolSize):
        threading.Thread.__init__(self, name = "Scheduler", target = self.Run)
        self.pool = ThreadPool(poolSize)
        self.tasks = Queue.Queue()

    def ExecuteOne(self, blocking = True):
        logging.debug("Looking for next task...")
        try:
            task = self.tasks.get(blocking)
        except Queue.Empty:
            logging.debug("No task to run")
            return None
        result = None
        error = None
        traceback = None
        while True:
            cb = task._GetNext()
            if not cb:
                # no more callback
                break
            if cb.threaded:
                # a threaded callback
                self._AddJob(task, cb, result, error, traceback)
                # don't pass Go, don't reclaim $200
                return None
            # Run the callback according to the current state
            try:
                if error:
                    error = cb.errorback(error)
                else:
                    result = cb.callback(result)
            except:
                errtype, error, traceback = sys.exc_info()
        if error:
            raise error, None, traceback
        else:
            return result

    def Run(self):
        logging.info("Scheduler start")
        while True:
            try:
                self.ExecuteOne()
            except self._SchedulerStop:
                break
            except:
                logging.exception("Unhandled task exception")
        logging.info("Scheduler stop")

    def Start(self):
        self.pool.Start()
        return self.start()

    def Stop(self, now = False):
        self.pool.Stop(now)
        if now:
            self.tasks = Queue.Queue()
        # We raise an exception to find if we stop stop the scheduler.
        # We could have use a None task, but this make it easier if we
        # want to add such mechanism public or we want to stop on
        # other exception
        def RaiseSchedulerStop():
            raise self._SchedulerStop
        self.AddTask(Task(RaiseSchedulerStop))
        self.join()

    def AddTask(self, task):
        self.tasks.put(task)

    def _AddJob(self, task, cb, result, error, traceback):

        def DoIt(task, cb, result, error, traceback):
            try:
                if error:
                    error = cb.errorback(error)
                else:
                    result = cb.callback(result)
            except:
                errtype, error, traceback = sys.exc_info()
            if error:
                def RaiseError():
                    raise error, None, traceback
                jobTask = Task(RaiseError)
            else:
                def ReturnResult():
                    return result
                jobTask = Task(ReturnResult)
            jobTask.ChainTask(task)
            self.AddTask(jobTask)

        # This double wrap (Job over DoIt) seems necessary to make
        # error not look like a local of Job...
        def Job():
            return DoIt(task, cb, result, error, traceback)
        self.pool.AddJob(Job)

# The global scheduler
scheduler = None

def StartScheduler(size):
    global scheduler
    if scheduler:
        StopScheduler()
    scheduler = Scheduler(size)
    scheduler.Start()

def StopScheduler(now = False):
    global scheduler
    if scheduler:
        scheduler.Stop(now)
    scheduler = None

if __name__ == '__main__':
    from time import sleep
    logging.getLogger().setLevel(logging.INFO)
    if len(sys.argv) == 1:
    	numberOfTasks = 100
    else:
        numberOfTasks = int(sys.argv[1])
    # This function is a sample and shouldn't know about the scheduler
    count = 0
    def AsyncCall(name, seconds):
        global count
        count += 1
        
        # Probably a bad example, since the callback
        # doesn't return the exact same type...
        def Initialize(name, seconds):
            print "Here", name
            return name, seconds
        def Blocking(args):
            name, time = args
            print name, "goes to bed"
            sleep(time)
            print name, ": ZZZ..."
            return name
        def Finalize(name):
            global count
            print name, "wakes up!"
            count -= 1
            return name

        task = Task(Initialize, name, seconds)
        task.AddThreadedCallback(Blocking)
        task.AddCallback(Finalize)
        return task

    logging.info("Starting scheduler with 10 workers")
    StartScheduler(10)
    logging.info("Adding asynccall task")
    for x in xrange(numberOfTasks):
        task = AsyncCall("Toto%d" % (x+1), (x % 10)/10.0)
        scheduler.AddTask(task)
    while count > 0:
        logging.debug("Count = %d", count)
        sleep(1)

    # Check for King Toto sleep
    task = AsyncCall("King Toto", 5)
    data = task.GrabResult()
    event = task.SetupEvent()
    scheduler.AddTask(task)
    try:
        event.wait(10)
        print "data = %r" % (data,)
    except:
        logging.exception("Error occured on wait")
    logging.info("Stopping scheduler")
    StopScheduler()
    logging.info("The End.")