# -*- coding: utf-8 -*-

import os
import logging
import contextlib

import yt.wrapper as yt


class SaaSYtCypress(object):
    """
    Main class for work with yt cypress.
    Wiki: https://wiki.yandex-team.ru/yt/userdoc/pythonwrapper/
    """
    def __init__(self, proxy='', basedir='', token=''):
        if not token:
            token = os.environ.get('YT_TOKEN')
        else:
            logging.warning('Yt token was not found!')
        self._yt = yt.YtClient(proxy=proxy, token=token)
        self._basedir = basedir
        self._current = self._basedir + '/current'
        self._finish = self._basedir + '/finish'

    # Base actions
    def __create(self, path, ntype='map_node'):
        result = ''
        try:
            result = self._yt.create(ntype, path)
        except Exception:
            logging.exception('Yt create error')
        if result:
            return True
        else:
            return False

    def __get(self, path):
        dirs = dict()
        try:
            dirs = self._yt.get(path)
        except Exception:
            logging.exception('Yt path error')
        return dirs

    def __list(self, path):
        try:
            self._yt.list(path)
        except Exception:
            logging.exception('Yt list error')

    def __set(self, path, data):
        try:
            self._yt.set(path, data)
        except Exception:
            logging.exception('Yt set error')

    def __move(self, srcpath, dstpath):
        try:
            self._yt.move(srcpath, dstpath, force=True)
        except Exception:
            logging.exception('Yt move error')

    def __remove(self, path):
        try:
            self._yt.remove(path)
        except Exception:
            logging.exception('Yt remove error')

    def _get_jobs_list(self):
        return self.__get(self._current).keys()

    def _get_job_attrs(self, job):
        return self.__get(self._current + '/' + job + '/@attrs')

    def _get_job_locks(self, job):
        return self.__get(self._current + '/' + job + '/@locks')

    def _set_job_attrs(self, job, attrs):
        self.__set(self._current + '/' + job + '/@attrs', attrs)

    def _check_for_jobs(self):
        return self.__get(self._current)

    def _create_job(self, job, attrs):
        job_request = self.__create(self._current + '/' + job)
        if job_request:
            logging.info('Job created successfully')
            self._set_job_attrs(job, attrs)
        else:
            logging.error('Failed to create job')

    def _remove_lock(self, job):
        locks = self._get_job_locks(job)
        if len(locks):
            for lock in locks:
                transaction_id = lock['transaction_id']
                logging.info('Removing lock transaction %s for job %s', transaction_id, job)
                self._yt.abort_transaction(transaction_id)
        else:
            logging.info('Locks for job %s was not found', job)
        return self.__get(self._current + '/' + job)

    def _remove_job(self, job):
        self.__remove(self._current + '/' + job)

    def _finish_job(self, job):
        self.__move(self._current + '/' + job, self._finish + '/' + job)

    def _unfinish_job(self, job):
        self.__move(self._finish + '/' + job, self._current + '/' + job)


class SaaSYtJobControl(SaaSYtCypress):
    JOB_STAGES = ['job', 'service', 'delivery']
    JOB_STAGE_STATUSES = ['NEW', 'IN_PROGRESS', 'FINISH']

    def __init__(self, proxy='hahn', basedir='//home/saas/ssm', token=''):
        super(SaaSYtJobControl, self).__init__(proxy=proxy, basedir=basedir, token=token)

    def new_service_job(self, service_name, service_ctype, service_type, instances_count, req_memory='', sla_info='', delivery_info='', dc='',
                        prepare_env=False, service_shard_by='', replicas_per_dc='1', startrek_issue=False, req_cpu='2', template_service='',
                        no_indexing=False, allocation_type='yp', service_tvm='', service_saas_tvm=False):
        """
        Creates new service request. All types - str or dict
        :param service_name: type str
        :param service_ctype: type str
        :param service_type: type str
        :param instances_count: type str
        :param req_memory: type str
        :param sla_info: type dict
        :param delivery_info: type dict
        :param dc: type str
        :param prepare_env: type boolean
        :param service_shard_by: type str
        :param replicas_per_dc: type str
        :param startrek_issue: type boolean
        :param req_cpu: type str
        :param no_indexing: type boolean
        :return:
        """
        service_job_name = service_ctype + '-' + service_name
        service_attrs = {
            'job_type': 'service',
            'allocation_type': allocation_type,
            'service_name': service_name,
            'service_ctype': service_ctype,
            'service_type': service_type,
            'service_shard_by': service_shard_by,
            'instances_count': str(instances_count),
            'dc': dc,
            'req_memory': str(req_memory),
            'req_cpu': str(req_cpu),
            'sla_info': sla_info,
            'delivery_info': delivery_info,
            'no_indexing': no_indexing,
            'prepare_env': prepare_env,
            'template_service': template_service,
            'startrek_issue': startrek_issue,
            'replicas_per_dc': str(replicas_per_dc),
            'service_tvm': service_tvm,
            'service_saas_tvm': service_saas_tvm,
            'status': {
                'job': 'NEW',
                'service': 'NEW',
                'delivery': 'NEW'
            },
            'approved': False
        }
        self._create_job(service_job_name, service_attrs)

    def new_namespace_job(self, namespace_name, ferryman_service, ferryman_ctype, owners='', size='', doccount='', ticket=''):
        """
        Creates new namespace request
        :param namespace_name: type str
        :param ferryman_service: type str
        :param ferryman_ctype: type str
        :param owners: type str or list
        :param size: type int or str
        :param doccount: type int or str
        :param ticket: type str
        """
        namespace_job_name = 'ns-' + ferryman_service + '-' + namespace_name
        service_attrs = {
            'job_type': 'namespace',
            'namespace_name': namespace_name,
            'ferryman_service': ferryman_service,
            'ferryman_ctype': ferryman_ctype,
            'owners_list': owners,
            'namespace_size': str(size),
            'namespace_doccount': str(doccount),
            'sla_info': {
                'ticket': ticket
            },
            'status': {
                'job': 'NEW',
                'service': 'NEW',
            },
            'approved': False
        }
        self._create_job(namespace_job_name, service_attrs)

    def get_approve_status(self, job):
        """
        Get job approve status. Available values - True, False.
        :param job: type str
        :return: type boolean
        """
        job_attrs = self._get_job_attrs(job)
        return job_attrs['approved']

    def get_avail_jobs(self):
        """
        Get list of non-finished jobs.
        :return: type list
        """
        return self._get_jobs_list()

    def get_job_info(self, job):
        """
        Get information about job
        :param job: type str
        :return: type dict
        """
        return self._get_job_attrs(job)

    def get_job_locks(self, job):
        """
        Get list of locks for job
        :param job: type str
        :return: type list
        """
        return self._get_job_locks(job)

    def discard_request(self, job):
        """
        Discard request, remove yt map node.
        :param job: type str
        :return:
        """
        return self._remove_job(job)

    def set_approve_status(self, job, approve):
        """
        Change job approve status. Available values - True, False.
        :param job: type str
        :param approve: type boolean
        :return:
        """
        if approve in [True, False]:
            job_attrs = self._get_job_attrs(job)
            job_attrs['approved'] = approve
            self._set_job_attrs(job, job_attrs)

    def check_job(self, job, stage='job'):
        """
        Check job for approve, status and locks
        :param job: type str
        :param stage: type str in ['job', 'service', 'delivery']
        :return: type boolean
        """
        job_attrs = self._get_job_attrs(job)
        job_status = job_attrs['status'][stage]
        job_approved_state = job_attrs.get('approved')
        if job_status not in ['STOP', 'FINISH'] and bool(job_approved_state) and len(self.get_job_locks(job)) == 0:
            return True
        else:
            return False

    def check_job_exists(self, job_name):
        jobs = self.get_avail_jobs()
        if job_name in jobs:
            return True
        else:
            return False

    def take_job(self, job):
        """
        Change job status to IN_PROCESS
        :param job: type str
        :return:
        """
        job_attrs = self._get_job_attrs(job)
        job_attrs['status']['job'] = 'IN_PROCESS'
        self._set_job_attrs(job, job_attrs)

    def finish_job(self, job):
        """
        Change job status to FINISH and move job from /current to /finish.
        :param job: type str
        :return:
        """
        job_attrs = self._get_job_attrs(job)
        if job_attrs['status']['service'] == 'FINISH' and (job_attrs['status'].get('delivery') == 'FINISH' or job_attrs['job_type'] == 'namespace'):
            self.change_job_status(job, 'job', 'FINISH')
            self._finish_job(job)
        else:
            logging.warning('Try to finish job is FAILED because one of stages not in status FINISH')

    def finish_service_stage(self, job):
        """
        Change status of job service stage to FINISH
        :param job: type str
        :return:
        """
        self.change_job_status(job, 'service', 'FINISH')

    def finish_delivery_stage(self, job):
        """
        Change status of job delivery stage to FINISH
        :param job: type str
        :return:
        """
        self.change_job_status(job, 'delivery', 'FINISH')

    def start_service_stage(self, job):
        """
        Change status of job service stage to IN_PROCESS
        :param job: type str
        :return:
        """
        self.change_job_status(job, 'service', 'IN_PROCESS')

    def start_delivery_stage(self, job):
        """
        Change status of job delivery stage to IN_PROCESS
        :param job: type str
        :return:
        """
        self.change_job_status(job, 'delivery', 'IN_PROCESS')

    def change_job_status(self, job, stage, status):
        """
        Change job status
        :param job:
        :param stage: type str in ['job', 'service', 'delivery']
        :param status: type str in ['NEW', 'IN_PROCESS', 'FINISH']
        :return:
        """
        if stage not in ['job', 'service', 'delivery']:
            logging.error('Stage is not valid')
        if status not in ['NEW', 'IN_PROCESS', 'STOP', 'FINISH']:
            logging.error('Status is not valid')
        job_attrs = self._get_job_attrs(job)
        job_attrs['status'][stage] = status
        self._set_job_attrs(job, job_attrs)

    def reborn_job(self, job, service_stage=True, delivery_stage=True):
        logging.info('Reborn job %s', job)
        self._unfinish_job(job)
        self.change_job_status(job, 'job', 'NEW')
        if service_stage:
            self.change_job_status(job, 'service', 'NEW')
        if delivery_stage:
            self.change_job_status(job, 'delivery', 'NEW')

    @contextlib.contextmanager
    def run_job(self, job, timeout=10):
        """
        Context manager for lock and run jobs.
        :param job: type str
        :param timeout: type int
        :return:
        """
        logging.info('Running job %s', job)
        self.take_job(job)
        # Running job
        try:
            with self._yt.Transaction():
                self._yt.lock(self._current + '/' + job, mode='exclusive', wait_for=timeout)
                yield
        finally:
            logging.info('Stopping job %s', job)
