# Copyright (c) 2020 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com

from enum import Enum
from datetime import datetime, timedelta, timezone
from typing import Mapping

import boto3
import json
import time

from czc.aws import get_paginated_results
from czc.exceptions import AlreadyTrackingExecutionException
from czc.logging import get_logger
from czc.step_functions import StateMachine

step_functions_client = boto3.client('stepfunctions')
logger = get_logger()


class Status(str, Enum):
    running = 'RUNNING'
    succeeded = 'SUCCEEDED'
    failed = 'FAILED'
    timed_out = 'TIMED_OUT'
    aborted = 'ABORTED'


detail_map = {
    'ActivityFailed': 'activityFailedEventDetails',
    'ActivityScheduled': 'activityScheduledEventDetails',
    'ActivityScheduleFailed': 'activityScheduleFailedEventDetails',
    'ActivityStarted': 'activityStartedEventDetails',
    'ActivitySucceeded': 'activitySucceededEventDetails',
    'ActivityTimedOut': 'activityTimedOutEventDetails',
    'ChoiceStateEntered': 'stateEnteredEventDetails',
    'ChoiceStateExited': 'stateExitedEventDetails',
    'ExecutionAborted': 'executionAbortedEventDetails',
    'ExecutionFailed': 'executionFailedEventDetails',
    'ExecutionStarted': 'executionStartedEventDetails',
    'ExecutionSucceeded': 'executionSucceededEventDetails',
    'ExecutionTimedOut': 'executionTimedOutEventDetails',
    'FailStateEntered': 'stateEnteredEventDetails',
    'LambdaFunctionFailed': 'lambdaFunctionFailedEventDetails',
    'LambdaFunctionScheduled': 'lambdaFunctionScheduledEventDetails',
    'LambdaFunctionScheduleFailed': 'lambdaFunctionScheduleFailedEventDetails',
    'LambdaFunctionStarted': None,
    'LambdaFunctionStartFailed': 'lambdaFunctionStartFailedEventDetails',
    'LambdaFunctionSucceeded': 'lambdaFunctionSucceededEventDetails',
    'LambdaFunctionTimedOut': 'lambdaFunctionTimedOutEventDetails',
    'MapIterationAborted': 'mapIterationAbortedEventDetails',
    'MapIterationFailed': 'mapIterationFailedEventDetails',
    'MapIterationStarted': 'mapIterationStartedEventDetails',
    'MapIterationSucceeded': 'mapIterationSucceededEventDetails',
    'MapStateAborted': None,
    'MapStateEntered': 'stateEnteredEventDetails',
    'MapStateExited': 'stateExitedEventDetails',
    'MapStateFailed': None,
    'MapStateStarted': 'mapStateStartedEventDetails',
    'MapStateSucceeded': None,
    'ParallelStateAborted': None,
    'ParallelStateEntered': 'stateEnteredEventDetails',
    'ParallelStateExited': 'stateExitedEventDetails',
    'ParallelStateFailed': None,
    'ParallelStateStarted': None,
    'ParallelStateSucceeded': None,
    'PassStateEntered': 'stateEnteredEventDetails',
    'PassStateExited': 'stateExitedEventDetails',
    'SucceedStateEntered': 'stateEnteredEventDetails',
    'SucceedStateExited': 'stateExitedEventDetails',
    'TaskFailed': 'taskFailedEventDetails',
    'TaskScheduled': 'taskScheduledEventDetails',
    'TaskStarted': 'taskStartedEventDetails',
    'TaskStartFailed': 'taskStartFailedEventDetails',
    'TaskStateAborted': None,
    'TaskStateEntered': 'stateEnteredEventDetails',
    'TaskStateExited': 'stateExitedEventDetails',
    'TaskSubmitFailed': 'taskSubmitFailedEventDetails',
    'TaskSubmitted': 'taskSubmittedEventDetails',
    'TaskSucceeded': 'taskSucceededEventDetails',
    'TaskTimedOut': 'taskTimedOutEventDetails',
    'WaitStateAborted': None,
    'WaitStateEntered': 'stateEnteredEventDetails',
    'WaitStateExited': 'stateExitedEventDetails'
}


class StandardStateMachine(StateMachine):
    '''
    This class should be used instead of `StateMachine`.
    '''
    def __init__(self, arn: str, base_name: str):
        super().__init__(arn, base_name)
        self._execution_arn = None

    def _convert_states(self, states):
        def get_detail_for_state(state):
            detail_field = detail_map.get(state['type'])
            return state.get(detail_field, {}) if detail_field else {}

        return [{
            'id': state.get('id'),
            'previousEventId': state.get('previousEventId'),
            'type': state.get('type'),
            'timestamp': state.get('timestamp'),
            'details': get_detail_for_state(state)
        } for state in states]

    def start(self, **kwargs):
        '''
        This method can be used to start an execution of the state machine.

        Any named arguments passed to this method will be added as arguments to pass to the state machine.
        Additionall, the arguments "execution_name" and "start_timestamp" will be added to the arguments passed
        to the state machine.
        '''
        if self._execution_arn:
            raise AlreadyTrackingExecutionException('Cannot start an execution - already tracking another execution.')

        response = step_functions_client.start_execution(
            stateMachineArn=self._arn,
            name=self.name,
            input=json.dumps(kwargs)
        )

        self._execution_arn = response['executionArn']
        return self._execution_arn

    def get_execution(self):
        if not self._execution_arn:
            raise Exception('Execution was never started.')

        try:
            return step_functions_client.describe_execution(executionArn=self._execution_arn)
        except Exception:
            return None

    def get_completion_status(self):
        execution = self.get_execution()
        if not execution or execution['status'] == Status.running:
            return None

        response = step_functions_client.get_execution_history(
            executionArn=self._execution_arn,
            maxResults=1,
            reverseOrder=True)

        return self._convert_states(response['events'])[0]

    def wait_for_complete(self, timeout: int = 300, sleep_interval: int = 5) -> Mapping:
        end_time = datetime.now(timezone.utc) + timedelta(seconds=timeout)

        while (status := self.get_completion_status()) is None:
            if datetime.now(timezone.utc) > end_time:
                raise TimeoutError()
            time.sleep(sleep_interval)

        return status

    def get_states(self):
        events = get_paginated_results(step_functions_client, 'get_execution_history', 'events', executionArn=self._execution_arn)
        return self._convert_states(events)
