import logging
import time

import boto3

import config

logging.basicConfig(
    format="%(asctime)s %(levelname)s %(message)s", level=config.LOG_LEVEL
)

stepfunctions_client = boto3.client("stepfunctions", region_name=config.AWS_REGION)


def start_state_machine(
    state_machine_arn, execution_name=None, state_machine_input=None
):
    start_execution_args = {"stateMachineArn": state_machine_arn}

    if execution_name:
        start_execution_args["name"] = execution_name

    if state_machine_input:
        start_execution_args["input"] = state_machine_input

    logging.info(f"Starting execution for state machine {state_machine_arn}")
    response = stepfunctions_client.start_execution(**start_execution_args)
    execution_arn = response["executionArn"]
    logging.info(f"Started execution with ARN {execution_arn}")
    return execution_arn


def wait_for_state_machine_to_complete(execution_arn, polling_interval):
    while True:
        response = stepfunctions_client.describe_execution(executionArn=execution_arn)
        status = response["status"]
        if status == "RUNNING":
            logging.info(
                f"State machine is running. Waiting for {polling_interval} seconds"
            )
            time.sleep(polling_interval)
        elif status == "SUCCEEDED":
            logging.info("State machine completed successfully.")
            break
        else:
            message = f"State machine did not complete successfully. Status: {status}."
            message += f'\nError: {response["error"]}' if response.get("error") else ""
            message += f'\nCause: {response["cause"]}' if response.get("cause") else ""
            raise Exception(message)


def main():
    """Execute main entrypoint."""
    state_machine_arn = config.STATE_MACHINE_ARN

    if not state_machine_arn:
        raise Exception("STATE_MACHINE_ARN is required")

    execution_arn = start_state_machine(
        state_machine_arn,
        config.STATE_MACHINE_EXECUTION_NAME,
        config.STATE_MACHINE_INPUT,
    )

    wait_for_state_machine_to_complete(
        execution_arn, config.STATE_MACHINE_POLLING_INTERVAL
    )


if __name__ == "__main__":
    main()