"""
Utilitários para interagir com o AWS Step Functions, incluindo iniciar execuções, aguardar a conclusão e gerenciar
payloads.

Este módulo fornece funções para iniciar execuções de máquinas de estado do Step Functions, aguardar a conclusão das
execuções e gerenciar payloads de entrada e saída. As funções são projetadas para serem usadas em scripts de CI/CD e
automação de tarefas relacionadas ao AWS Step Functions, garantindo que as execuções sejam realizadas de forma confiável
e que os resultados sejam validados conforme necessário.

"""

from __future__ import annotations

import json
import secrets
import string
import time
from pathlib import Path

from common import print_message
from src.datalake.commons.aws_utils import (
    make_client,
    make_session,
    resolve_stack_output,
)

WAIT_FOR_STEPFUNCTION_TIME = 30
WAIT_START_STEPFUNCTION_EXECUTION_TIME = 5


def _random_id(size: int = 10) -> str:
    """
    Gera um ID aleatório.

    Args:
        size (int): Tamanho do ID a ser gerado.

    Returns:
        str: ID aleatório gerado.

    """
    alphabet = string.ascii_letters + string.digits
    return "".join(secrets.choice(alphabet) for _ in range(size))


def _resolve_state_machine_arn(output_key: str, env: dict[str, str]) -> str:
    """
    Resolve o ARN da máquina de estado do Step Functions com base na chave de saída do stack.

    Args:
        output_key (str): Chave de saída do stack que contém o ARN da máquina de estado.
        env (dict[str, str]): Dicionário contendo as variáveis de ambiente.

    Returns:
        str: ARN da máquina de estado.

    """
    session = make_session(env)
    arn = resolve_stack_output(
        session,
        stack_name=env.get("STACK_NAME", ""),
        output_key=output_key,
        region_name=env.get("REGION", ""),
    )
    if not arn:
        raise RuntimeError(
            f"[Dados] stepfunctions - output não encontrado para chave: {output_key}"
        )
    return arn


def start_stepfunction_execution(
    stepfunction_name: str, payload: dict[str, str], env: dict[str, str]
) -> str:
    """
    Inicia a execução de uma máquina de estado do Step Functions.

    Args:
        stepfunction_name (str): Nome da máquina de estado.
        payload (dict[str, str]): Payload de entrada para a execução.
        env (dict[str, str]): Dicionário contendo as variáveis de ambiente.

    Returns:
        str: ARN da execução iniciada.

    """
    file = Path("response.json")
    print_message(
        f"[Dados] stepfunctions - start-execution - executando pipeline {stepfunction_name}"
    )

    arn = _resolve_state_machine_arn(stepfunction_name, env)
    print_message(f"Variavel {arn}")

    if "execution_name" not in payload:
        payload["execution_name"] = f"deploy-{Path.cwd().name}-{_random_id()}"

    file_payload = Path("payload.json")
    payload_str = json.dumps(payload)
    file_payload.write_text(payload_str, encoding="utf-8")

    print_message(f"[Dados] stepfunctions start-execution: {arn} {payload}")
    session = make_session(env)
    sf = make_client(session, "stepfunctions", region_name=env.get("REGION", ""))
    response = sf.start_execution(
        stateMachineArn=arn,
        input=payload_str,
    )
    file.write_text(json.dumps(response, default=str), encoding="utf-8")

    print_message("[Dados] stepfunctions start-execution response")
    print_message(file.read_text(encoding="utf-8"))

    execution_arn = str(response.get("executionArn", ""))
    if not execution_arn:
        raise RuntimeError("[Dados] Falha ao iniciar execução de stepfunction")
    print_message(f"executionArn: {execution_arn}")
    return execution_arn


def wait_for_stepfunction_execution(
    stepfunction_name: str, execution_arn: str, env: dict[str, str]
) -> None:
    """
    Aguarda a conclusão de uma execução de máquina de estado do Step Functions.

    Args:
        stepfunction_name (str): Nome da máquina de estado.
        execution_arn (str): ARN da execução da máquina de estado.
        env (dict[str, str]): Dicionário contendo as variáveis de ambiente.

    Returns:
        None

    """
    print_message(f"[Dados] stepfunctions - describe-execution: {stepfunction_name}")
    session = make_session(env)
    sf = make_client(session, "stepfunctions", region_name=env.get("REGION", ""))
    while True:
        response = sf.describe_execution(executionArn=execution_arn)
        content = json.dumps(response, default=str)
        print_message(content)

        status = str(response.get("status", ""))
        status_upper = status.upper()
        if status_upper in {"FAILED", "TIMED_OUT", "ABORTED"}:
            raise RuntimeError(
                f"[Dados] stepfunctions - Erro de execução encontrado, invalidando job! status={status_upper}"
            )
        if status_upper in {"RUNNING"}:
            print_message(
                f"[Dados] stepfunctions - {stepfunction_name} ainda em execução, tentando novamente em "
                f"{WAIT_FOR_STEPFUNCTION_TIME}s..."
            )
            time.sleep(WAIT_FOR_STEPFUNCTION_TIME)
            continue

        print_message(f"[Dados] stepfunctions - {stepfunction_name} ok!")
        return


def start_and_wait_stepfunctions_execution(
    names: list[str], payload: dict[str, str], env: dict[str, str]
) -> None:
    """
    Inicia e aguarda a execução de múltiplas máquinas de estado do Step Functions.

    Args:
        names (list[str]): Lista de nomes das máquinas de estado.
        payload (dict[str, str]): Payload de entrada para as execuções.
        env (dict[str, str]): Dicionário contendo as variáveis de ambiente.

    Returns:
        None

    """
    execution_map: list[tuple[str, str]] = []
    for stepfunction_name in names:
        execution_arn = start_stepfunction_execution(
            stepfunction_name, dict(payload), env
        )
        execution_map.append((stepfunction_name, execution_arn))

    print_message(
        f"[Dados] stepfunctions - Aguardando {WAIT_START_STEPFUNCTION_EXECUTION_TIME}..."
    )
    time.sleep(WAIT_START_STEPFUNCTION_EXECUTION_TIME)

    for stepfunction_name, execution_arn in execution_map:
        wait_for_stepfunction_execution(stepfunction_name, execution_arn, env)

    print_message("[Dados] stepfunctions - tudo ok!")