"""
Execução de modelos de ML no Data Lake.

Este módulo fornece funções para executar modelos de ML, incluindo a execução de máquinas de estado do Step Functions e
a invocação de funções Lambda. As funções são projetadas para serem usadas em scripts de CI/CD e automação de tarefas
relacionadas ao Data Lake, garantindo que os modelos de ML sejam executados de forma confiável e consistente.

"""

from __future__ import annotations

import argparse
import secrets
import string
from pathlib import Path

from common import print_message
from src.datalake.commons.data_env import load_runtime_env
from src.datalake.commons.lambda_utils import invoke_lambda
from src.datalake.commons.state_machine_utils import (
    start_and_wait_stepfunctions_execution,
)


def _random_id(size: int = 10) -> str:
    alphabet = string.ascii_letters + string.digits
    return "".join(secrets.choice(alphabet) for _ in range(size))


def run_ml(wait_finish: bool, state_machine_id: str, env: dict[str, str]) -> None:
    mode = "aguardando termino" if wait_finish else "sem aguardar termino"
    print_message(f"[Dados] executando modelo de ML - {mode}")

    payload = {
        "execution_name": f"deploy-{Path.cwd().name}-{_random_id()}",
        "predicate": "-",
    }

    start_and_wait_stepfunctions_execution([state_machine_id], payload, env)

    if wait_finish:
        invoke_lambda("InsertApiRequisitesFunctionId", None, env)


def main() -> int:
    parser = argparse.ArgumentParser(description="Data CI ML helpers")
    parser.add_argument("-w", action="store_true", dest="wait")
    parser.add_argument("-n", action="store_true", dest="no_wait")
    parser.add_argument("state_machine_id")
    args = parser.parse_args()

    try:
        env = load_runtime_env(apply_set_all=True, configure_aws=True)
        wait_finish = True
        if args.no_wait:
            wait_finish = False
        if args.wait:
            wait_finish = True
        run_ml(wait_finish, args.state_machine_id, env)
    except Exception as exc:
        print_message(str(exc))
        return 2

    return 0


if __name__ == "__main__":
    raise SystemExit(main())