from __future__ import with_statement
from functools import partial
import logging
import threading

from imvu.task import task, Future, Return, Start, Sleep, CompletedFuture, isSchedulable

logger = logging.getLogger('imvu.' + __name__)

def _assertFutures(futures):
    assert all(isinstance(f, Future) for f in futures)

@task
def WaitForAll(futures):
    futures = list(futures)
    _assertFutures(futures)
    
    for f in futures:
        yield f
    yield Return(futures)

class RaceJudge(object):
    def __init__(self):
        self.racedict = {}
    def tryToWin(self):
        bogus = object()
        return bogus == self.racedict.setdefault(1, bogus)

def WaitForFirst(futures):
    futures = list(futures)
    _assertFutures(futures)

    firstFuture = Future()

    raceJudge = RaceJudge()
    def onComplete(future, result, error):
        assert error is None
        if raceJudge.tryToWin():
            slow_futures = set(futures)
            fast_future = future
            slow_futures.remove(fast_future)
            result = (fast_future, slow_futures)
            firstFuture.complete(result, None)

    for f in futures:
        f.registerOnComplete(partial(onComplete, f))

    return firstFuture

@task
def RunInParallel(tasks):
    futures = [(yield Start(t)) for t in tasks]
    yield Return([(yield f) for f in futures])

class TaskTimeout(Exception):
    pass
Timeout = TaskTimeout

@task
def WaitWithTimeout(task, timeout):
    task_future    = yield Start(task)
    timeout_future = yield Start(Sleep(timeout))
    yield WaitForFirst([task_future, timeout_future])
    if task_future.isComplete():
        yield Return(task_future.result)
    else:
        raise Timeout

@task
def WaitForEvent(event, timeout):
    yield Return((yield WaitWithTimeout(event.wait(), timeout)))

class Event(object):
    def __init__(self):
        self.__lock = threading.Lock()
        self.__waitingFutures = []
        self.__isSet = False

    def wait(self):
        with self.__lock:
            if self.__isSet:
                return CompletedFuture(None)
            
            f = Future()
            self.__waitingFutures.append(f)
            return f

    def clear(self):
        with self.__lock:
            self.__isSet = False

    def set(self):
        with self.__lock:
            self.__isSet = True
            waiters = self.__waitingFutures
            self.__waitingFutures = []
        for f in waiters:
            f.complete(None, None)

    def pulse(self):
        self.set()
        self.clear()

    def isSet(self):
        return self.__isSet

class Cache(object):
    def __init__(self, task):
        self.__task = task
        self.__results = {}

    @task
    def __call__(self, *args):
        key = args
        if key not in self.__results:
            self.__results[key] = (yield self.__task(*args))
        yield Return(self.__results[key])

@task
def CallPeriodically(task_or_function, period):
    while True:
        r = task_or_function()
        if isSchedulable(r):
            yield r
            
        yield Sleep(period)

class Queue(object):
    def __init__(self):
        self.__q = []
        self.__waitingFutures = []
        self.__lock = threading.Lock()

    def empty(self):
        return not self.__q
        
    def put(self, item):
        with self.__lock:
            try:
                future = self.__waitingFutures.pop(0)
            except IndexError:
                self.__q.append(item)
                return

        # If we complete the future while the lock is held, we could
        # deadlock.  Arbitrary code can run in a future's set of
        # onCompletes.
        future.complete(item, None)
        
    def get(self):
        with self.__lock:
            try:
                item = self.__q.pop(0)
            except IndexError:
                f = Future()
                self.__waitingFutures.append(f)
                return f
        
        return CompletedFuture(item)