"""
Utilitários para executar e validar consultas no Redshift, incluindo execução de statements e verificação de resultados.

Este módulo fornece funções para iniciar statements no Redshift, aguardar a conclusão das consultas, verificar se há
erros ou duplicatas nos resultados e gerenciar o fluxo de execução de consultas em lote. As funções são projetadas para
serem usadas em scripts de CI/CD e automação de tarefas relacionadas ao Redshift, garantindo que as consultas sejam
executadas de forma confiável e que os resultados sejam validados conforme necessário.

"""

from __future__ import annotations

import json
import time
from pathlib import Path

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

WAIT_FOR_REDSHIFT_QUERY_TIME = 30
WAIT_START_REDSHIFT_QUERY_TIME = 5


def start_redshift_query(table_name: str, query_sql: str, env: dict[str, str]) -> str:
    """
    Inicia a execução de um statement no Redshift.

    Args:
        table_name (str): Nome da tabela para a qual o statement está sendo executado.
        query_sql (str): O statement SQL a ser executado.
        env (dict[str, str]): Dicionário contendo variáveis de ambiente necessárias para a execução do statement.

    Returns:
        str: O ID do statement no Redshift.

    """
    print_message(f"[Dados] redshift - execute-statement: {table_name}: '{query_sql}'")
    redshift_environment = (
        "dev" if env.get("ENVIRONMENT") == "hom" else env.get("ENVIRONMENT")
    )
    session = make_session(env)
    redshift_data = make_client(
        session, "redshift-data", region_name=env.get("REGION", "")
    )
    response = redshift_data.execute_statement(
        Sql=query_sql,
        DbUser=f"servico.dataoffice.flyway.{env.get('ENVIRONMENT')}",
        ClusterIdentifier=f"{redshift_environment}-general-data-platform",
        Database=env.get("ENVIRONMENT", ""),
    )
    response_text = json.dumps(response, default=str)
    Path("response.json").write_text(response_text, encoding="utf-8")
    print_message(response_text)
    query_id = str(response.get("Id", ""))
    if not query_id:
        raise RuntimeError("[Dados] Falha ao iniciar statement no Redshift")
    print_message(f"Id: {query_id}")
    return query_id


def wait_for_redshift_query(
    table_name: str, query_id: str, env: dict[str, str], verification: str
) -> tuple[bool, bool]:
    """
    Aguarda a conclusão de um statement no Redshift e verifica os resultados.

    Args:
        table_name (str): Nome da tabela para a qual o statement está sendo executado.
        query_id (str): O ID do statement no Redshift.
        env (dict[str, str]): Dicionário contendo variáveis de ambiente necessárias para a execução do statement.
        verification (str): Descrição da verificação que está sendo realizada.

    Returns:
        tuple[bool, bool]: Uma tupla indicando se houve erro e se houve duplicatas.

    """
    session = make_session(env)
    redshift_data = make_client(
        session, "redshift-data", region_name=env.get("REGION", "")
    )

    while True:
        body = redshift_data.describe_statement(Id=query_id)
        print_message(json.dumps(body, default=str))
        status = str(body.get("Status", "")).upper()

        if status in {"FAILED", "ABORTED"}:
            print_message(
                "[Dados] redshift - Erro de execução encontrado, invalidando CI!"
            )
            return True, False

        if status in {"STARTED", "SUBMITTED", "PICKED"}:
            print_message(
                f"[Dados] redshift - {table_name} ainda em execução, tentando novamente em "
                f"{WAIT_FOR_REDSHIFT_QUERY_TIME}s..."
            )
            time.sleep(WAIT_FOR_REDSHIFT_QUERY_TIME)
            continue

        rows = int(body.get("ResultRows") or 0)
        print_message(f"[Dados] redshift - {table_name} qtd: {rows}")
        if rows >= 1:
            print_message(
                f"[Dados] redshift - {verification} encontrado, invalidando CI!"
            )
            return False, True

        print_message(f"[Dados] redshift - {table_name} ok!")
        return False, False


def start_and_wait_redshift_queries(
    pairs: list[tuple[str, str]], env: dict[str, str], verification: str
) -> None:
    """
    Inicia e aguarda a execução de múltiplos statements no Redshift.

    Args:
        pairs (list[tuple[str, str]]): Lista de tuplas contendo o nome da tabela e o statement SQL.
        env (dict[str, str]): Dicionário contendo variáveis de ambiente necessárias para a execução dos statements.
        verification (str): Descrição da verificação que está sendo realizada.

    Returns:
        None

    """
    query_ids: list[tuple[str, str]] = []
    for table_name, query_sql in pairs:
        query_id = start_redshift_query(table_name, query_sql, env)
        query_ids.append((table_name, query_id))

    print_message(f"[Dados] redshift - Aguardando {WAIT_START_REDSHIFT_QUERY_TIME}...")
    time.sleep(WAIT_START_REDSHIFT_QUERY_TIME)

    errors: list[str] = []
    duplicates: list[str] = []
    for table_name, query_id in query_ids:
        has_error, has_duplicate = wait_for_redshift_query(
            table_name, query_id, env, verification
        )
        if has_error:
            errors.append(table_name)
        if has_duplicate:
            duplicates.append(table_name)

    if errors:
        raise RuntimeError(
            f"[Dados] Erro de execução, invalidando CI! Encontrado erro de execução em: {', '.join(errors)}"
        )
    if duplicates:
        raise RuntimeError(
            f"[Dados] Falha: {verification}, invalidando CI! Dados {verification} em: {', '.join(duplicates)}"
        )

    print_message("[Dados] redshift ok!")