# Copyright (C) 2005-2024 Splunk Inc. All Rights Reserved.

import itsi_py3
import os
from ITOA.itoa_object import ItoaObject
from ITOA.setup_logging import logger
from ITOA.itoa_common import get_current_utc_epoch


class ItsiUpgradeReadinessPrechecks(ItoaObject):
    """
    Implements ITSI Upgrade Readiness Precheck object
    """
    logger = logger
    log_prefix = '[ITSI Upgrade Readiness Precheck] '
    collection_name = 'itsi_upgrade_readiness_prechecks'
    object_type = 'upgrade_readiness_prechecks'

    def __init__(self, session_key, current_user_name):
        session_key = session_key
        super(ItsiUpgradeReadinessPrechecks, self).__init__(session_key,
                                                            current_user_name,
                                                            self.object_type,
                                                            collection_name=self.collection_name,
                                                            title_validation_required=True)

    """
    Miscellaneous helper methods
    """

    def create_upgrade_readiness_precheck_job(self, transaction_id, operation_mode):
        '''
        Creates a record of the upgrade readiness precheck job with the transaction id.

        @rtype: ItoaObject???
        @return: UpgradeReadinessPrecheck
        '''

        time_now = get_current_utc_epoch()

        new_upgrade_readiness_prechecks = []
        new_upgrade_readiness_prechecks.append(
            {
                'status': 'IN_PROGRESS',
                'start_time': time_now,
                'transaction_id': transaction_id,
                'title': f'upgrade_readiness_precheck_{transaction_id}',
                'operation_mode': operation_mode,
                'precheck_started': False,
                'remediation_started': False,
                'process_id': os.getpid()
            }
        )

        self.save_batch(
            'nobody', new_upgrade_readiness_prechecks, True)

    def update_upgrade_readiness_precheck_job(self, transaction_id, **kwargs):
        '''
        Updates the record of the upgrade readiness precheck job of the given transaction id with the new data.

        @rtype: ItoaObject
        @return: UpgradeReadinessPrecheck
        '''

        filter_data = {
            '$or': [{'$and': [
                {'transaction_id': transaction_id}]
            }]
        }

        upgrade_readiness_precheck_job = self.get_bulk(
            'nobody',
            fields=['_key', 'start_time', 'transaction_id',
                    'status', 'object_type', 'title', 'operation_mode', 'precheck_started',
                    'precheck_start_time', 'remediation_started', 'remediation_start_time', 'process_id'],
            filter_data=filter_data)

        updated_upgrade_readiness_precheck_job = []
        for job in upgrade_readiness_precheck_job:
            for key, value in kwargs.items():
                job[key] = value
            updated_upgrade_readiness_precheck_job.append(job)

        if updated_upgrade_readiness_precheck_job:
            self.save_batch(
                'nobody', updated_upgrade_readiness_precheck_job, True)

    def check_process(self, process_id):
        """Check whether the process is alive or not based on process ID. Returns True if job is alive else return False.
            : param process_id: process ID of existing precheck job.
        """
        try:
            os.kill(process_id, 0)
            return True
        except (OSError, TypeError):
            return False

    def get_in_progress_upgrade_readiness_prechecks(self, lookback_time=None, exclude_transaction_id=None, mode=None):
        """
        Gets prechecks from KVStore that are in progress.

        @type lookback_time: int
        @param lookback_time: range of time before current_time in seconds to query for precheck jobs

        @type exclude_transaction_id: string
        @param exclude_transaction_id: the transaction_id to exclude in the query

        @rtype: array
        @return: returns array of in progress jobs
        """

        filter_array = [{'status': 'IN_PROGRESS'}]
        if exclude_transaction_id:
            filter_array.append({'transaction_id': {'$ne': exclude_transaction_id}})

        if lookback_time:
            time_now = get_current_utc_epoch()
            earliest_time_stamp = time_now - lookback_time
            filter_array.append({'start_time': {'$gte': earliest_time_stamp}})

        if mode:
            filter_array.append({'operation_mode': mode})

        filter_data = {'$and': filter_array}

        other_in_progress_jobs = self.get_bulk(
            'nobody',
            filter_data=filter_data,
            fields=['_key', 'start_time', 'transaction_id',
                    'status', 'object_type', 'title', 'operation_mode', 'process_id'],
            sort_dir='desc',
            sort_key='start_time')

        if len(other_in_progress_jobs) > 0:
            updated_other_in_progress_jobs = []
            for in_progress_job in other_in_progress_jobs:
                if self.check_process(in_progress_job.get('process_id', None)):
                    updated_other_in_progress_jobs.append(in_progress_job)
                else:
                    self.update_upgrade_readiness_precheck_job(in_progress_job['transaction_id'], status='FAILED')
            other_in_progress_jobs = updated_other_in_progress_jobs

        return other_in_progress_jobs