from crawler.q import RedisQueue
from settings import REDIS_KEY


class Scheduler:

    def __init__(self, queue=None):
        self.mrq = (queue or RedisQueue())

    def get_task(self, redis_key=None):
        if redis_key is not None:
            return self.mrq.pop_task_by_key(redis_key)
        else:
            return self.mrq.pop_task([REDIS_KEY])

    def insert_task(self, task, level=1, **kwargs):
        """插入一条任务"""
        self.mrq.push_task(REDIS_KEY, [task], level=level, **kwargs)

    def insert_tasks(self, tasks, level=1, **kwargs):
        """批量插入任务"""
        allow_output_log = kwargs.get('allow_output_log', True)
        self.mrq.push_task(
            REDIS_KEY,
            tasks,
            level=level,
            allow_output_log=allow_output_log
        )