import threading
from google.protobuf import message as _message


class TaskTable(object):
    def __init__(self):
        self.taskMap = {}
        self.lock = threading.Lock()

    def PushTask(self, primary_key, msg: _message.Message):
        if len(self.taskMap.items()) > 0:
            return False
        else:
            self.lock.acquire()
            self.taskMap[primary_key] = msg
            self.lock.release()
            return True

    def UpdateResult(self, primary_key, ret):
        if primary_key in self.taskMap:
            taskTable = self.taskMap[primary_key]
            self.lock.acquire()
            self.taskMap[primary_key] = [taskTable, ret]
            self.lock.release()
            return True

        else:
            return False

    def GetTask(self):
        if len(self.taskMap.items()) == 0:
            return None
        else:
            for item in self.taskMap.items():
                key, task = item
                if isinstance(task, (_message.Message)):
                    return task
                else:
                    return None

    def GetResult(self, primary_key):
        if len(self.taskMap.items()) == 0:
            return None
        else:
            if primary_key in self.taskMap:
                result = self.taskMap[primary_key]
                if isinstance(result, (list)):
                    ret = result[1]
                else:
                    return None
                self.lock.acquire()
                self.taskMap.clear()
                self.lock.release()
                return ret
            else:
                return None