
from .models import Session, Trigger, Spider, Project, SpiderExecutionQueue, HistoricalJob, session_scope, \
    SpiderSettings, Node
from apscheduler.schedulers.tornado import TornadoScheduler
from apscheduler.executors.pool import ThreadPoolExecutor, ProcessPoolExecutor
from apscheduler.triggers.cron import CronTrigger
from sqlite3 import IntegrityError
from tornado.ioloop import IOLoop, PeriodicCallback
import uuid
import logging
import datetime
from .exceptions import *
from .config import Config
from sqlalchemy import distinct, desc, or_, and_
import os
from .mail import MailSender
from six import string_types



def generate_job_id():
    jobid = uuid.uuid4().hex
    return jobid

logger = logging.getLogger(__name__)

JOB_STATUS_PENDING = 0
JOB_STATUS_RUNNING = 1
JOB_STATUS_SUCCESS = 2
JOB_STATUS_FAIL = 3
JOB_STATUS_WARNING = 4
JOB_STATUS_STOPPING = 5


class SchedulerManager():
    def __init__(self, config=None, syncobj=None):
        if config is None:
            config = Config()
        self.config = config
        executors = {
            'default': ThreadPoolExecutor(20),
            'processpool': ProcessPoolExecutor(5)
        }
        self.scheduler = TornadoScheduler(executors=executors)

        self.poll_task_queue_callback = None
        self.pool_task_queue_interval = 10
        self.ioloop = IOLoop.instance()
        self.clear_finished_jobs_callback = PeriodicCallback(self.clear_finished_jobs, 60*1000)
        self.reset_timeout_job_callback = PeriodicCallback(self.reset_timeout_job, 10*1000)

        self.sync_obj = syncobj
        if syncobj is not None:
            self.sync_obj.set_on_remove_schedule_job(self.on_cluster_remove_scheduling_job)
            self.sync_obj.set_on_add_schedule_job(self.on_cluster_add_scheduling_job)

    def init(self):
        session = Session()

        # move completed jobs into history
        for job in session.query(SpiderExecutionQueue).filter(SpiderExecutionQueue.status.in_((2, 3))):
            historical_job = HistoricalJob()
            historical_job.id = job.id
            historical_job.spider_id = job.spider_id
            historical_job.project_name = job.project_name
            historical_job.spider_name = job.spider_name
            historical_job.fire_time = job.fire_time
            historical_job.start_time = job.start_time
            historical_job.complete_time = job.update_time
            historical_job.status = job.status
            session.delete(job)
            session.add(historical_job)
        session.commit()

        # init triggers
        triggers = session.query(Trigger)
        for trigger in triggers:
            try:
                self.add_job(trigger.id, trigger.cron_pattern)
            except InvalidCronExpression:
                logger.warning('Trigger %d,%s cannot be added ' % (trigger.id, trigger.cron_pattern))
        session.close()

        self.scheduler.start()

        self.clear_finished_jobs_callback.start()
        self.reset_timeout_job_callback.start()

    def build_cron_trigger(self, cron):
        cron_parts = cron.split(' ')
        if len(cron_parts) != 5:
            raise InvalidCronExpression()

        try:
            crontrigger = CronTrigger(minute=cron_parts[0],
                                      hour=cron_parts[1],
                                      day=cron_parts[2],
                                      month=cron_parts[3],
                                      day_of_week=cron_parts[4],
                                      )
            return crontrigger
        except ValueError:
            raise InvalidCronExpression()

    def add_job(self, trigger_id, cron):
        logger.debug('adding trigger %s %s' % (trigger_id, cron))
        crontrigger = self.build_cron_trigger(cron)


        job = self.scheduler.add_job(func=self.trigger_fired, trigger=crontrigger, kwargs={'trigger_id': trigger_id},
                               id=str(trigger_id), replace_existing=True)
        if self.sync_obj:
            self.ioloop.call_later(0, self.sync_obj.add_schedule_job, trigger_id)
            #self.sync_obj.add_schedule_job(trigger_id)

    def on_cluster_remove_scheduling_job(self, job_id):
        logger.debug('on_cluster_remove_scheduling_job')
        if self.scheduler.get_job(job_id):
            self.scheduler.remove_job(job_id)

    def on_cluster_add_scheduling_job(self, trigger_id):
        logger.debug('on_cluster_add_scheduling_job')
        with session_scope() as session:
            trigger = session.query(Trigger).filter_by(id = trigger_id).first()
            if trigger is None:
                return
            crontrigger = self.build_cron_trigger(trigger.cron_pattern)
            job = self.scheduler.add_job(func=self.trigger_fired, trigger=crontrigger,
                                         kwargs={'trigger_id': trigger_id},
                                         id=str(trigger_id), replace_existing=True)

    def trigger_fired(self, trigger_id):
        with session_scope() as session:
            trigger = session.query(Trigger).filter_by(id=trigger_id).first()
            if not trigger:
                logger.error('Trigger %s not found.' % trigger_id)
                return

            spider = session.query(Spider).filter_by(id=trigger.spider_id).first()
            if not spider:
                logger.error('Spider %s not found' % spider.name)
                return

            project = session.query(Project).filter_by(id=spider.project_id).first()
            if not project:
                logger.error('Project %s not found' % project.name)
                return

        try:
            self.add_task(project.name, spider.name)
        except JobRunning:
            logger.info('Job for spider %s.%s already reach the concurrency limit' % (project.name, spider.name))


    def add_schedule(self, project, spider, cron):
        with session_scope()as session:
            triggers = session.query(Trigger).filter(Trigger.spider_id==spider.id)
            found = False
            for trigger in triggers:
                if trigger.cron_pattern == cron:
                    found = True
                    break

            if not found:
                # create a cron_trigger for just validating
                cron_trigger = self.build_cron_trigger(cron)
                trigger = Trigger()
                trigger.spider_id = spider.id
                trigger.cron_pattern = cron
                session.add(trigger)
                session.commit()
                self.add_job(trigger.id, cron)

    def add_task(self, project_name, spider_name):
        with session_scope() as session:
            project = session.query(Project).filter(Project.name==project_name).first()
            spider = session.query(Spider).filter(Spider.name==spider_name, Spider.project_id==project.id).first()

            executing = list(session.query(SpiderExecutionQueue).filter(SpiderExecutionQueue.spider_id==spider.id, SpiderExecutionQueue.status.in_([0,1])))
            concurrency_setting = session.query(SpiderSettings).filter_by(spider_id=spider.id, setting_key='concurrency').first()
            concurrency = int(concurrency_setting.value) if concurrency_setting else 1
            executing_slots = [executing_job.slot for executing_job in executing]
            free_slots = [x for x in range(1, concurrency + 1) if x not in executing_slots]
            if not free_slots:
                logger.warning('spider %s-%s is configured as %d concurency, and %d in queue, skipping' % (
                    project.name, spider.name, concurrency, len(executing_slots)))
                raise JobRunning(executing[0].id)

            executing = SpiderExecutionQueue()
            spider_tag_vo = session.query(SpiderSettings).filter_by(spider_id=spider.id, setting_key='tag').first()
            spider_tag = spider_tag_vo.value if spider_tag_vo else None
            jobid = generate_job_id()
            executing.id = jobid
            executing.spider_id = spider.id
            executing.project_name = project.name
            executing.spider_name = spider.name
            executing.fire_time = datetime.datetime.now()
            executing.update_time = datetime.datetime.now()
            executing.tag = spider_tag
            executing.slot = free_slots[0]
            session.add(executing)
            session.commit()
            session.refresh(executing)
            return executing

    def on_node_expired(self, node_id):
        session = Session()
        for job in session.query(SpiderExecutionQueue).filter(SpiderExecutionQueue.node_id==node_id, SpiderExecutionQueue.status == 1):
            job.status = 0
            job.update_time = datetime.datetime.now()
            job.start_time = None
            job.pid = None
            job.node_id = None
            session.add(job)
        session.commit()
        session.close()

    def jobs(self):
        session = Session()
        pending = list(session.query(SpiderExecutionQueue).filter(SpiderExecutionQueue.status==0))
        running = list(session.query(SpiderExecutionQueue).filter(SpiderExecutionQueue.status==1))
        finished = list(session.query(HistoricalJob).order_by(desc(HistoricalJob.complete_time)).slice(0, 100))
        session.close()
        return pending, running, finished

    def job_start(self, jobid, pid):
        with session_scope() as session:
            job = session.query(SpiderExecutionQueue).filter_by(id=jobid).first()
            if job.start_time is None:
                job.start_time = datetime.datetime.now()
            job.update_time = datetime.datetime.now()
            if job.pid is None and pid:
                job.pid = pid
            session.add(job)
            session.commit()
            session.close()

    def _regular_agent_tags(self, agent_tags):
        if agent_tags is None:
            return None
        if isinstance(agent_tags, string_types):
            return agent_tags.split(',')
        return agent_tags

    def get_next_task(self, node_id):
        with session_scope() as session:
            node = session.query(Node).filter(Node.id == node_id).first()
            node_tags = node.tags
            next_task = self._get_next_task(session, node_tags)
            if not next_task:
                return None
            next_task.start_time = datetime.datetime.now()
            next_task.update_time = datetime.datetime.now()
            next_task.node_id = node_id
            next_task.status = 1
            session.add(next_task)
            session.commit()
            session.refresh(next_task)
            return next_task
        return None

    def get_next_task_of_tags(self, tags, node_id):
        with session_scope() as session:
            next_task = self._get_next_task(session, tags)
            if not next_task:
                return None
            next_task.start_time = datetime.datetime.now()
            next_task.update_time = datetime.datetime.now()
            next_task.node_id = node_id
            next_task.status = 1
            session.add(next_task)
            session.commit()
            session.refresh(next_task)
            return next_task
        return None

    def _get_next_task(self, session, agent_tags):
        agent_tags = self._regular_agent_tags(agent_tags)
        tag_settings = session.query(SpiderSettings).filter(SpiderSettings.setting_key == 'tag').subquery()
        query = session.query(SpiderExecutionQueue)\
            .join(SpiderExecutionQueue.spider)\
            .join(tag_settings, Spider.settings, isouter=True).filter(SpiderExecutionQueue.status == 0)

        if agent_tags:
            tags_condition = tag_settings.c.value.in_(agent_tags)
        else:
            tags_condition = tag_settings.c.value.is_(None)


        query = query.filter(tags_condition)
        query = query.order_by(SpiderExecutionQueue.fire_time)
        next_task = query.first()
        if next_task:
            return next_task

        return None

    def has_task(self, node_id):
        with session_scope() as session:
            node = session.query(Node).filter(Node.id == node_id).first()
            if node is None:
                return None
            node_tags = node.tags
            next_task = self._get_next_task(session, node_tags)
        return next_task is not None


    def jobs_running(self, node_id, job_ids):
        '''

        :param node_id:
        :param job_ids:
        :return:(job_id) to kill
        '''
        with session_scope() as session:
            for job_id in job_ids:
                job = session.query(SpiderExecutionQueue).filter(
                    SpiderExecutionQueue.id == job_id).first()

                if job:
                    if job.node_id is None:
                        job.node_id = node_id
                    if job.node_id != node_id or \
                        job.status != 1:
                        yield job.id
                    else:
                        job.update_time = datetime.datetime.now()
                        session.add(job)
                else:
                    yield job_id
            session.commit()

    def job_finished(self, job, log_file=None, items_file=None):
        session = Session()
        if job.status not in (2,3):
            raise Exception('Invliad status.')
        job_status = job.status
        job = session.query(SpiderExecutionQueue).filter_by(id=job.id).first()
        job.status = job_status
        job.update_time = datetime.datetime.now()

        historical_job = HistoricalJob()
        historical_job.id = job.id
        historical_job.spider_id = job.spider_id
        historical_job.project_name = job.project_name
        historical_job.spider_name = job.spider_name
        historical_job.fire_time = job.fire_time
        historical_job.start_time = job.start_time
        historical_job.complete_time = job.update_time
        historical_job.status = job.status
        if log_file:
            historical_job.log_file = log_file
            import re
            items_crawled_pattern = re.compile("\'item_scraped_count\': (\d+),")
            error_log_pattern = re.compile("\'log_count/ERROR\': (\d+),")
            warning_log_pattern = re.compile("\'log_count/WARNING\': (\d+),")
            with open(log_file, 'r') as f:
                log_content = f.read()
                m = items_crawled_pattern.search(log_content)
                if m:
                    historical_job.items_count = int(m.group(1))

                m = error_log_pattern.search(log_content)
                if m and historical_job.status == JOB_STATUS_SUCCESS:
                    historical_job.status = JOB_STATUS_FAIL
                m = warning_log_pattern.search(log_content)
                if m and historical_job.status == JOB_STATUS_SUCCESS:
                    historical_job.status = JOB_STATUS_WARNING

        if items_file:
            historical_job.items_file = items_file
        session.delete(job)
        session.add(historical_job)
        session.commit()
        session.refresh(historical_job)

        # send mail
        if historical_job.status == JOB_STATUS_FAIL:
            self.try_send_job_failed_mail(historical_job)


        session.close()
        return historical_job

    def try_send_job_failed_mail(self, job):
        logger.debug('try_send_job_failed_mail')
        job_fail_send_mail = self.config.getboolean('job_fail_send_mail')
        if job_fail_send_mail:
            try:
                mail_sender = MailSender(self.config)
                subject = 'scrapydd job failed'
                to_address = self.config.get('job_fail_mail_receiver')
                content = 'bot:%s \r\nspider:%s \r\n job_id:%s \r\n' % (job.spider.project.name,
                                                                        job.spider_name,
                                                                        job.id)
                mail_sender.send(to_addresses=to_address, subject=subject, content=content)

            except Exception as e:
                logger.error('Error when sending job_fail mail %s' % e)

    def clear_finished_jobs(self):
        job_history_limit_each_spider = 100
        with session_scope() as session:
            spiders = list(session.query(distinct(HistoricalJob.spider_id)))
        for row in spiders:
            spider_id = row[0]
            with session_scope() as session:
                over_limitation_jobs = list(session.query(HistoricalJob)\
                    .filter(HistoricalJob.spider_id==spider_id)\
                    .order_by(desc(HistoricalJob.complete_time))\
                    .slice(job_history_limit_each_spider, 1000)\
                    .all())
            for over_limitation_job in over_limitation_jobs:
                self._remove_histical_job(over_limitation_job)

    def _clear_running_jobs(self):
        with session_scope() as session:
            jobs = list(session.query(SpiderExecutionQueue).filter(SpiderExecutionQueue.status.in_([0,1])))
        for job in jobs:
            self._remove_histical_job(job)

    def reset_timeout_job(self):
        KILL_TIMEOUT = 120
        with session_scope() as session:
            timeout_time = datetime.datetime.now() - datetime.timedelta(minutes=1)
            for job in session.query(SpiderExecutionQueue).filter(SpiderExecutionQueue.status == 1):
                spider = session.query(Spider).filter_by(id=job.spider_id).first()
                job_timeout_setting = session.query(SpiderSettings).filter_by(spider_id=spider.id,
                                                                              setting_key='timeout').first()
                job_timeout = int(job_timeout_setting.value) if job_timeout_setting else 3600
                logger.debug((job.update_time - job.start_time).seconds)
                if job.update_time < timeout_time:
                    # job is not refresh as expected, node might be died, reset the status to PENDING
                    job.status = 0
                    job.pid = None
                    job.node_id = None
                    job.update_time = datetime.datetime.now()
                    session.add(job)
                    logger.info('Job %s is timeout, reseting.' % job.id)
                elif (job.update_time - job.start_time).seconds > job_timeout:
                    # let agent kill the job and continue submit a log.
                    job.status = JOB_STATUS_STOPPING
                    session.add(job)
            for job in session.query(SpiderExecutionQueue).filter(SpiderExecutionQueue.status.in_([1,5])):
                spider = session.query(Spider).filter_by(id=job.spider_id).first()
                job_timeout_setting = session.query(SpiderSettings).filter_by(spider_id=spider.id,
                                                                              setting_key='timeout').first()
                job_timeout = int(job_timeout_setting.value) if job_timeout_setting else 3600
                if (job.update_time - job.start_time).seconds > job_timeout + KILL_TIMEOUT:

                    # job is running too long, should be killed
                    historical_job = HistoricalJob()
                    historical_job.id = job.id
                    historical_job.spider_id = job.spider_id
                    historical_job.project_name = job.project_name
                    historical_job.spider_name = job.spider_name
                    historical_job.fire_time = job.fire_time
                    historical_job.start_time = job.start_time
                    historical_job.complete_time = job.update_time
                    historical_job.status = 3
                    session.delete(job)
                    session.add(historical_job)
                    logger.info('Job %s is timeout, killed.' % job.id)
            session.commit()

    def _remove_histical_job(self, job):
        '''
        @type job: HistoricalJob
        '''

        with session_scope() as session:
            job = session.query(HistoricalJob).filter(HistoricalJob.id == job.id).first()
            if job.items_file:
                try:
                    os.remove(job.items_file)
                except Exception as e:
                    logger.warning(e.message)

            if job.log_file:
                try:
                    os.remove(job.log_file)
                except Exception as e:
                    logger.warning(e.message)

            original_log_file = os.path.join('logs', job.project_name, job.spider_name, '%s.log' % job.id)
            if os.path.exists(original_log_file):
                os.remove(original_log_file)

            original_items_file = os.path.join('items', job.project_name, job.spider_name, '%s.jl' % job.id)
            if os.path.exists(original_items_file):
                os.remove(original_items_file)
            session.delete(job)

    def remove_schedule(self, project_name, spider_name, trigger_id):
        with session_scope() as session:
            project = session.query(Project).filter(Project.name == project_name).first()
            spider = session.query(Spider).filter(Spider.project_id == project.id, Spider.name == spider_name).first()
            trigger = session.query(Trigger).filter(Trigger.spider_id==spider.id, Trigger.id == trigger_id).first()

            session.delete(trigger)
            if self.scheduler.get_job(str(trigger_id)):
                self.scheduler.remove_job(str(trigger.id))

        if self.sync_obj:
            logger.info('remove_schedule')
            self.sync_obj.remove_schedule_job(trigger.id)
