#!/usr/bin/env python3
# coding: utf-8

"""

"""
import sys
import threading
from pymongo import MongoClient
from datetime import datetime, timedelta
import argparse
import platform
import subprocess
import json
from stocks3.core.config import Configuration


def pprint(obj):
    if '_id' in obj:
        obj['_id'] = str(obj['_id'])
    print(json.dumps(obj, indent=4))


# Индексы для базы:
# db.planner.createIndex({"interval.from": 1, "interval.to": 1}, {background: true})
# храним записи за три дня:
# db.planner_log.createIndex({ts: 1}, {background: true, expireAfterSeconds: 86400*3})
# db.planner_log.createIndex({job_id: 1}, {background: true})
# db.planner_log.createIndex({host: 1}, {background: true})
# db.planner_lock.createIndex({"job_id": 1}, {background: true})
# db.planner_lock.createIndex({"host": 1}, {background: true})
# db.planner_lock.createIndex({"last_start": 1}, {background: true})


def do_job(job, planner_object):
    """
    Выполняет задание в отдельном процессе и записывает вывод в job_output.
    По завершении задания выполняет release_job
    :param job:
    :param planner_object:
    :return:
    """
    try:
        print('doing:', job['command'])
        output = subprocess.check_output(job['command'], shell=True, stderr=subprocess.STDOUT)
        job_output = output.decode('utf-8').strip()
        error = False
        print('finished OK:', job['command'])
    except subprocess.CalledProcessError as e:
        job_output = e.output.decode('utf-8').strip()
        error = True
        print('finished ERROR:', job['command'])

    planner_object.release_job(job, job_output, error)


class Planner(object):
    def __init__(self, connection_string, database, max_lock_time, max_bind_time):
        self.connection_string = connection_string
        self.database = database
        self.max_lock_time = max_lock_time
        self.max_bind_time = max_bind_time
        self.client = None
        self.db = None
        self.planner_config = None
        self.planner_lock = None
        self.planner_log = None
        self.host = platform.node()

        self.connect()

    def connect(self):
        self.client = MongoClient(self.connection_string)
        self.db = self.client[self.database]
        self.planner_config = self.db.planner
        self.planner_log = self.db.planner_log
        self.planner_lock = self.db.planner_lock

    def acquire_job(self, job, lock):
        """
        Пытаемся заблокировать задачу для ее выполнения
        :param job: dict
        :param lock: dict
        :return: boolean
        """
        find_conditions = {'job_id': job['_id']}
        if job['distribution'] == 'all':
            find_conditions['host'] = self.host

        if lock is None or lock['host'] != self.host:
            # upsert нужен для того, чтобы не появилось две записи в БД
            update_set = {
                'job_id': job['_id'],
                'host': self.host
            }
            print('upserting lock record:', self.planner_lock.update(find_conditions, update_set, upsert=True))
        elif 'acquired' in lock and 'ts' in lock['acquired']:
            if datetime.now() - lock['acquired']['ts'] > timedelta(seconds=self.max_lock_time):
                r = self.planner_lock.update({'_id': lock['_id']}, {'$unset': {'acquired': 1}})
                print('found stale job! releasing:', r)

        find_conditions['acquired'] = {'$exists': False}
        result = self.planner_lock.find_and_modify(
            query=find_conditions,
            update={'$set': {'acquired': {'host': self.host, 'ts': datetime.now()}, 'last_start': datetime.now()}},
            full_response=True)
        return result['ok'] and result['value'] is not None

    def release_job(self, job, output, error=False):
        """
        ПО завершению выполнения задачи убираем блокировку в БД
        :param job: dict
        :param output: str
        :param error: boolean
        :return: None
        """
        # print ('trying release job:', job)
        if job['distribution'] == 'one':
            self.planner_lock.update({'job_id': job['_id']},
                                     {'$unset': {'acquired': 1}, '$set': {'last_end': datetime.now()}})
        else:
            self.planner_lock.update({'job_id': job['_id'], 'host': self.host},
                                     {'$unset': {'acquired': 1}, '$set': {'last_end': datetime.now()}})
        log_entry = {'job_id': job['_id'],
                     'host': self.host,
                     'ts': datetime.now(),
                     'output': output}
        if error:
            log_entry['error'] = error
        self.planner_log.insert(log_entry)

    def check_last_start(self, job, now=None):
        """
        Проверка интервала между текущим временем и временем последнего запуска задачи.
        Возвращает False если запуск задачи возможен
        :param job:
        :param now: datetime
        :return: boolean
        """
        if now is None:
            now = datetime.now() + timedelta(seconds=5)
        period_text = tuple(int(component) for component in job['interval']['period'].split(':'))
        hours, minutes = period_text if len(period_text) > 1 else [0] + list(period_text)
        period = timedelta(hours=hours, minutes=minutes)

        if job["distribution"] == 'one':
            job_start = self.planner_lock.find_one({'job_id': job['_id']})
        else:
            job_start = self.planner_lock.find_one({'job_id': job['_id'], 'host': self.host})

        if job_start is None:
            print('job_start is None')
            return False, job_start

        if job.get('binding', 0) == 1:
            if self.host != job_start['host']:
                return (job_start['last_start'] + period + timedelta(seconds=self.max_bind_time) > now), job_start

        return job_start['last_start'] + period > now, job_start

    def check(self, now=datetime.now()):
        """
        Ищем в конфиге задачки которые должны выполняться в это время.
        :param now:datetime
        :return: None
        """
        if not self.planner_config:
            return
        string_time = now.strftime("%H:%M")
        contidions = {"$and": [{"interval.from": {"$lte": string_time}}, {"interval.to": {"$gte": string_time}}]}
        planner_records = self.planner_config.find(contidions)
        threads = []
        for job in planner_records:
            skip, lock = self.check_last_start(job, now=now)
            if skip:
                continue
            result = self.acquire_job(job, lock)
            if result:
                job_thread = threading.Thread(target=do_job, args=(job, self))
                job_thread.start()
                threads.append(job_thread)

        for job_thread in threads:
            job_thread.join()

    def list(self):
        """
        Показыаем список задач в конфиге
        :return: [str]
        """
        planner_records = self.planner_config.find({})
        for job in planner_records:
            job['text_interval'] = '{from}-{to} every {period}'.format(**job['interval'])
            if 'comment' not in job:
                job['comment'] = ''
            yield "Description: {comment}\nwhen: {text_interval}\ncmd: {command}\n".format(**job)

    def listrunning(self):
        """
        Показываем список выполняющихся сейчас задач
        :return: [str]
        """
        planner_records = self.planner_lock.find({'acquired': {'$exists': True}})
        for job in planner_records:
            yield job


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--database", help="Server connection string to connect", required=False)
    parser.add_argument("-c", "--config", help="Configuration file (default - config.json)", required=False)
    parser.add_argument("-l", "--list", help="List planner jobs", action='store_true', required=False)
    parser.add_argument("-r", "--running", help="List running jobs", action='store_true', required=False)
    args = parser.parse_args()

    config = Configuration(args.config) if args.config else Configuration()

    planner = Planner(config["db"]["connectionstring"],
                      config["db"]["planner_database"],
                      config['planner']['limits']['max_lock_time'],
                      config['planner']['limits']['max_lock_time'])
    if args.list:
        for job in planner.list():
            # print("\t".join((job['job_id'], job['host'], job['acquired'])))
            print("\t".join([job[k] for k in ('job_id', 'host', 'acquired')]))
        sys.exit(0)

    if args.running:
        for job in planner.listrunning():
            print(job['host'])
            # print("\t".join((job['job_id'], job['host'], job['acquired'])))
        # [print(job) for job in planner.listrunning()]
        sys.exit(0)

    planner.check(now=datetime.now())
