#!/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)

    # 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")

        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.")