"""
Executa validações customizadas de SQL no Data Lake usando Athena.

Este módulo fornece funções para executar queries SQL de validação no Athena, aguardando pelos resultados e gerenciando
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 Data Lake, garantindo que as validações sejam realizadas de forma confiável e que os resultados sejam
validados conforme necessário. Requer a existência da pasta ´sql/validation/athena´ com SQLs no projeto.

"""

from __future__ import annotations

import argparse
from pathlib import Path

from common import print_message
from src.datalake.commons.athena_utils import start_and_wait_athena_queries
from src.datalake.commons.ci_utils import extract_queries_from_file
from src.datalake.commons.data_env import load_runtime_env


def run_athena_sql_validations(data_type: str, env: dict[str, str]) -> None:
    env["ATHENA_VERIFICATION"] = "inconsistências em validações customizadas"
    print_message(
        f"[Dados] athena custom invoke - contando {env['ATHENA_VERIFICATION']} em todas as tabelas"
    )

    validation_folder = Path("sql/validation/athena")
    if not validation_folder.exists():
        print_message(
            f"[Dados] athena custom invoke - Diretório {validation_folder} não existe, seguindo..."
        )
        return

    all_queries: list[tuple[str, str]] = []
    for sql_file in validation_folder.rglob("*.sql"):
        print_message(f"[Dados] athena invoke - executando {sql_file} - {data_type}")
        all_queries.extend(extract_queries_from_file(env, sql_file, data_type))

    if all_queries:
        start_and_wait_athena_queries(all_queries, env, env["ATHENA_VERIFICATION"])


def main() -> int:
    parser = argparse.ArgumentParser(description="Data CI redshift athena validations")
    parser.add_argument("data_type")
    args = parser.parse_args()

    try:
        env = load_runtime_env(apply_set_all=True, configure_aws=True)
        run_athena_sql_validations(args.data_type, env)
    except Exception as exc:
        print_message(str(exc))
        return 2
    return 0


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