# 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 time
from collections.abc import Mapping
from datetime import datetime, timedelta, timezone
from typing import Optional

import boto3
from toolz.dicttoolz import get_in

import czc.cloudwatch as cloudwatch
from czc import json
from czc.date import get_current_timestamp
from czc.environ import namespace_aware
from czc.exceptions import AlreadyTrackingExecutionException, CannotFindLogGroupException, StepFunctionsNotFoundException
from czc.logging import get_logger

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

logger = get_logger()


@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'
        # Back track start time because we may not be in sync with AWS
        return self._convert_id(cloudwatch.query(self._log_group, query_string, self._epoch - 600))

    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'
            )
            # Back track start time because we may not be in sync with AWS
            logs = cloudwatch.query(self._log_group, query_string, self._epoch - 600, limit=max_states)
            records = self._convert_id([json.loads(item['@message']) for item in logs])

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

            time.sleep(sleep_interval)
