# 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

import json
import logging
import time

from pyfaaster.aws.handlers_decorators_v2 import namespace_aware
from datetime import datetime, timedelta, timezone
from typing import List, Mapping, Optional
from toolz.dicttoolz import get_in

import boto3

import czc.cloudwatch as cloudwatch
from czc.constants import PROJECT_NAME
from czc.date import get_current_timestamp
from czc.exceptions import AlreadyTrackingExecutionException, CannotFindLogGroupException, StepFunctionsNotFoundException


step_functions_client = boto3.client('stepfunctions')
tagging_api_client = boto3.client('resourcegroupstaggingapi')

logger = logging.getLogger(PROJECT_NAME)


@namespace_aware
def get_stm_arn_from_logical_id_tag(sfn_aws_logical_id: str, NAMESPACE):
    resources = tagging_api_client.get_resources(TagFilters=[
        {
            'Key': 'aws:cloudformation:logical-id',
            'Values': [sfn_aws_logical_id]
        },
        {
            'Key': 'cz:namespace',
            'Values': [NAMESPACE]
        }
    ])

    stm = next(filter(lambda r: r['ResourceARN'].startswith('arn:aws:states:'), resources['ResourceTagMappingList']), None)

    if not stm:
        raise StepFunctionsNotFoundException(f'Could not find state machine with tag aws:cloudformation:logical-id={sfn_aws_logical_id}')

    return stm['ResourceARN']


EXECUTION_STARTED = 'ExecutionStarted'
EXECUTION_SUCCEEDED = 'ExecutionSucceeded'
EXECUTION_FAILED = 'ExecutionFailed'


class ExpressStateMachine():
    _arn: str
    _base_name: str
    _epoch: int
    _log_group: str
    _execution_arn: Optional[str]
    _execution_name: Optional[str]

    @property
    def name(self):
        return self._execution_name

    def __init__(self, arn: str, base_name: str):
        self._arn = arn
        self._base_name = base_name
        self._epoch = int(get_current_timestamp())
        self._execution_name = f'{self._base_name}.{self._epoch}'
        self._execution_arn = None

        stm_info = step_functions_client.describe_state_machine(stateMachineArn=self._arn)
        log_destinations = get_in(['loggingConfiguration', 'destinations'], stm_info)

        if log_destinations is None or (log_destination := next(filter(lambda d: 'cloudWatchLogsLogGroup' in d, stm_info['loggingConfiguration']['destinations']), None)) is None:
            raise CannotFindLogGroupException(f'Cannot determine the log group for the state machine {self._arn}. Cannot track an express state machine without logs.')

        self._log_group = str(log_destination['cloudWatchLogsLogGroup']['logGroupArn']).split(':')[6]

    def _convert_id(self, state_logs: List[Mapping]) -> List[Mapping]:
        return [{
            **item,
            **({'id': int(item['id'])} if 'id' in item else {})
        } for item in state_logs]

    def _query_states(self, states: List[str]) -> List[Mapping]:
        states_filter = ' or '.join([f'type="{state}"' for state in states])
        query_string = f'fields @timestamp, execution_arn, type, id | filter execution_arn="{self._execution_arn}" and ({states_filter}) | sort @timestamp desc'
        return self._convert_id(cloudwatch.query(self._log_group, query_string, self._epoch - 600))  # Back track start time because we may not be in sync with AWS

    def get_completion_status(self):
        results = self._query_states([EXECUTION_FAILED, EXECUTION_SUCCEEDED])
        return results[0] if len(results) > 0 else None

    def track_execution(self, execution_arn: str, execution_name: str, start_time: int):
        '''This method can be used to attach this class to an execution that was already started.'''
        if self._execution_arn and self._execution_arn != execution_arn:
            raise AlreadyTrackingExecutionException(f'Cannot track execution {execution_arn} - already tracking execution {self._execution_arn}.')

        self._execution_arn = execution_arn
        self._execution_name = execution_name
        self._epoch = min(int(start_time), self._epoch)

    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.')

        args = {
            **kwargs,
            'execution_name': self._execution_name,
            'start_timestamp': self._epoch
        }

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

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

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

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

        return result

    def get_states(self, max_states: int = 5000, get_all: bool = False, timeout: int = 300, sleep_interval: int = 5):
        end_time = datetime.now(timezone.utc) + timedelta(seconds=timeout)
        last_id = None

        if get_all:
            completion_state = self.wait_for_complete(timeout, sleep_interval)
            last_id = completion_state['id']

        while True:
            query_string = f'fields @timestamp, @message, execution_arn, id | filter execution_arn="{self._execution_arn}" | sort id'
            logs = cloudwatch.query(self._log_group, query_string, self._epoch - 600, limit=max_states)  # Back track start time because we may not be in sync with AWS
            records = self._convert_id([json.loads(item['@message']) for item in logs])

            if last_id is None or len(records) == last_id:
                return records
            elif datetime.now(timezone.utc) > end_time:
                raise TimeoutError()

            time.sleep(sleep_interval)
