"""Step Functions and State Machine classes, with corresponding helper for SFN Integration tests."""
from __future__ import annotations

import time
from dataclasses import dataclass
from enum import Enum
from typing import List

import boto3
from botocore.exceptions import ClientError
from lambdacommon.common_config import logger
from lambdacommon.common_config import SFN_TEST_UNIQUE_IDENTIFIER
from lambdacommon.common_config import AWS_DEFAULT_REGION

SFN_CLIENT = boto3.client('stepfunctions', region_name=AWS_DEFAULT_REGION)

# To use COMMIT SHA for SFN_TEST_UNIQUE_IDENTIFIER in Integration Tests
# added truncated time to the minute, to ensure runs triggered around the same minute don't
# step on each other, but allowing for later retries on the same commit
UNIQUE_IDENTIFIER = f'{int(time.time() // 60 * 60)}-{SFN_TEST_UNIQUE_IDENTIFIER}'


class ExecutionStatus(str, Enum):
    """Execution status."""

    RUNNING = 'RUNNING'
    SUCCEEDED = 'SUCCEEDED'
    FAILED = 'FAILED'
    TIMED_OUT = 'TIMED_OUT'
    ABORTED = 'ABORTED'


class StateMachine:
    """Encapsulates SFN State Machine actions."""

    def __init__(self, arn: str):
        """Init a StateMachine."""
        self._arn = arn
        self._name = arn.split(':')[-1]

    @property
    def arn(self) -> str:
        """Get the State Machine ARN."""
        return self._arn

    @property
    def name(self) -> str:
        """Get the State Machine name."""
        return self._name

    def describe(self) -> dict:
        """Get information about the State Machine.

        Returns:
            The retrieved State Machine information.
        """
        try:
            return SFN_CLIENT.describe_state_machine(stateMachineArn=self.arn)
        except ClientError as err:
            logger.error(
                "Couldn't describe state machine %s. Here's why: %s: %s", self.arn,
                err.response['Error']['Code'], err.response['Error']['Message'])
            raise

    def start(self, execution_name: str, execution_input: str) -> str:
        """Start an execution of the State Machine with a specified name and input.

        Args:
            execution_name: Unique name for the execution.
            execution_input: Input for the State Machine execution.

        Returns:
            Execution ARN.
        """
        try:
            response = SFN_CLIENT.start_execution(
                stateMachineArn=self.arn,
                name=execution_name,
                input=execution_input,
            )
            return response['executionArn']
        except ClientError as err:
            logger.error(
                "Couldn't start state machine %s. Here's why: %s: %s", self.arn,
                err.response['Error']['Code'], err.response['Error']['Message'])
            raise

    def find_execution_by_name_prefix(
            self, name_prefix: str, status: ExecutionStatus = None) -> str | None:
        """Find an execution by a matching name prefix.

        Args:
            name_prefix: Execution name prefix to match.
            status: Optional execution status to filter by.

        Returns:
            Execution ARN found or None.
        """
        list_params = {}
        if status:
            list_params['statusFilter'] = status.value
        try:
            response = SFN_CLIENT.list_executions(stateMachineArn=self.arn, **list_params)
        except ClientError as err:
            logger.error("Couldn't list executions for state machine %s. Reason: %s: %s",
                         self.arn, err.response['Error']['Code'], err.response['Error']['Message'])
        else:
            for execution in response['executions']:
                if execution['name'].startswith(name_prefix):
                    return execution['executionArn']
            return None

    def describe_execution(self, execution_arn: str) -> dict:
        """Get execution information for a State Machine run such as current status or final output.

        Args:
            execution_arn: Execution ARN.

        Returns:
            Retrieved execution information.
        """
        try:
            return SFN_CLIENT.describe_execution(executionArn=execution_arn)
        except ClientError as err:
            logger.error(
                "Couldn't describe execution %s. Here's why: %s: %s", execution_arn,
                err.response['Error']['Code'], err.response['Error']['Message'])
            raise

    def wait_execution_end(self, execution_arn: str, polling_interval: int = 1) -> dict:
        """Wait for an execution to end and return execution info.

        Args:
            execution_arn: Execution ARN.
            polling_interval: Time in seconds to wait for polling execution status.

        Returns:
            Finished execution run information.
        """
        execution = self.describe_execution(execution_arn)
        while execution['status'] == ExecutionStatus.RUNNING:
            logger.info(f'Execution {execution_arn} still running, waiting for completion...')
            time.sleep(polling_interval)
            execution = self.describe_execution(execution_arn)
        return execution

    def get_all_task_history(
        self, execution_arn: str, include_execution_data: bool = True, results_per_page: int = 120
    ) -> List[dict]:
        """Retrieve all execution history task events for an execution.

        Args:
            execution_arn: Execution ARN.
            include_execution_data: Include event execution data like input and output.
            results_per_page: Amount of results to request per API call.

        Returns:
            List of all the execution history task events.
        """
        execution_history = []
        logger.info("Retrieving task's full execution history.")
        response = SFN_CLIENT.get_execution_history(
            executionArn=execution_arn,
            maxResults=results_per_page,
            includeExecutionData=include_execution_data
        )
        execution_history += response['events']
        # nextToken only included in the response if there is more events to retrieve
        while response.get('nextToken'):
            response = SFN_CLIENT.get_execution_history(
                executionArn=execution_arn,
                maxResults=results_per_page,
                includeExecutionData=include_execution_data,
                nextToken=response['nextToken']
            )
            execution_history += response['events']
        return execution_history


class StateType(str, Enum):
    """Execution status."""

    ENTERED = 'ENTERED'
    EXITED = 'EXITED'
    SCHEDULED = 'SCHEDULED'
    STARTED = 'STARTED'
    SUCCEEDED = 'SUCCEEDED'


class ExecutedState:
    """Representation of a State Machine executed State from the execution history."""

    def __init__(self, event):
        """Init an ExecutedState."""
        self._id = int(event['id'])
        self._previous_id = int(event['previousEventId'])
        self._state_type = event['type']
        for k, v in event.items():
            if k.endswith('EventDetails'):
                self._state_name = v.get('name')
                self._details = v

    @property
    def id(self) -> int:  # noqa
        """Get the State ID."""
        return self._id

    @property
    def previous_id(self) -> int:
        """Get the previous State ID."""
        return self._previous_id

    @property
    def state_type(self) -> str:
        """Get the State Type."""
        return self._state_type

    @property
    def state_name(self) -> str | None:
        """Get the State Name."""
        return getattr(self, '_state_name', None)

    @property
    def details(self) -> dict | None:
        """Get the State Details."""
        return getattr(self, '_details', None)

    @property
    def simplified_type(self) -> StateType | None:
        """Simplify State Type to only include the type status."""
        if self.state_type.endswith('Entered'):
            return StateType.ENTERED
        elif self.state_type.endswith('Exited'):
            return StateType.EXITED
        elif self.state_type.endswith('Scheduled'):
            return StateType.SCHEDULED
        elif self.state_type.endswith('Succeeded'):
            return StateType.SUCCEEDED
        elif self.state_type.endswith('Started'):
            return StateType.STARTED
        else:
            return None

    def __repr__(self): # noqa
        return str(self.__dict__)

    def __eq__(self, other): # noqa
        return repr(self) == repr(other)


@dataclass
class SFNExecutionError:
    """Representation of a State Machine execution error."""

    error: str
    cause: str


class StateMachineExecutionTest:
    """Encapsulates a Step Functions state machine execution test."""

    def __init__(self, sfn_arn: str, test_name: str, execution_input: str):
        """Init a StateMachineExecutionTest."""
        self._state_machine = StateMachine(sfn_arn)
        self._test_name = test_name
        self._execution_input = execution_input
        # This would try to ensure that for one test case, the execution unique name is unique for:
        # the same commit SHA, and a run within the same minute, for the same test name.
        # Max length allowed is 80 characters
        self._execution_name = f'{test_name}-{UNIQUE_IDENTIFIER}'[:80]
        self._execution_arn = self._state_machine.find_execution_by_name_prefix(
            self._execution_name)
        self._was_started_parallel = self._execution_arn is not None
        self._completed_execution = None
        self._execution_history: List[ExecutedState] = []

    @property
    def was_started_parallel(self):
        """Check if the test was started in a simultaneous run."""
        return self._was_started_parallel

    def run(self):
        """Start a Test Run or wait for an existing one with the same identifier to finish."""
        self._execution_arn = self._state_machine.start(self._execution_name, self._execution_input)

    def wait_execution_end(self):
        """Wait for the test end."""
        self._completed_execution = self._state_machine.wait_execution_end(self._execution_arn)

    @property
    def final_status(self) -> str | None:
        """Get Test Run final status."""
        if self._completed_execution:
            return self._completed_execution['status']
        return None

    @property
    def final_output(self) -> str | None:
        """Get Test Run final output."""
        if self._completed_execution:
            return self._completed_execution.get('output')
        return None

    @property
    def execution_error(self) -> SFNExecutionError | None:
        """Get Test Run error and cause if there was any."""
        if self._completed_execution and self._completed_execution.get('error'):
            return SFNExecutionError(
                self._completed_execution.get('error'),
                self._completed_execution.get('cause')
            )
        return None

    @property
    def was_successful(self) -> bool:
        """Check Test Run finished successful."""
        if self.final_status:
            return self.final_status == ExecutionStatus.SUCCEEDED
        return False

    @property
    def was_failed(self) -> bool:
        """Check Test Run finished on a failed state."""
        if self.final_status:
            return self.final_status == ExecutionStatus.FAILED
        return False

    def refresh_execution_history(self) -> List[ExecutedState]:
        """Get list of executed tasks.

        Returns:
            List of ExecutedState.
        """
        history = self._state_machine.get_all_task_history(self._execution_arn)
        self._execution_history = [ExecutedState(event) for event in history]
        return self._execution_history

    @property
    def execution_history(self) -> List[ExecutedState]:
        """Get list of executed tasks.

        Returns:
            List of ExecutedState.
        """
        if not self._execution_history:
            self.refresh_execution_history()
        return self._execution_history

    def get_task(self, task_name: str, state_type: StateType = None) -> ExecutedState | None:
        """Get a task from the execution history by its name if it ran.

        Args:
            task_name: Name of the task to look up.
            state_type: Type of the task to look up.

        Returns:
            ExecutedState if found, None if not.
        """
        if state_type:
            for task in self.execution_history:
                if task.state_name == task_name and task.simplified_type == state_type:
                    return task
        else:
            for task in self.execution_history:
                if task.state_name == task_name:
                    return task
        return None

    def task_was_called(self, task_name: str) -> bool:
        """Check if a task occurred in the Test Run.

        Args:
            task_name: Name of the task to look up.

        Returns:
            True if the task is found, False if not.
        """
        if self.get_task(task_name):
            return True
        return False