"""
Utilitários para interagir com o Athena, incluindo execução de queries e espera por resultados.
"""

from __future__ import annotations

import json
import time
from pathlib import Path
from urllib.parse import urlparse

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

WAIT_FOR_ATHENA_QUERY_TIME = 30
WAIT_START_ATHENA_QUERY_TIME = 5


def start_athena_query(table_name: str, query_sql: str, env: dict[str, str]) -> str:
    """
    Inicia a execução de uma query no Athena.

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

    Returns:
        str: O ID da execução da query no Athena.

    """
    print_message(
        f"[Dados] athena - start-query-execution: {table_name}: '{query_sql}'"
    )
    session = make_session(env)
    athena = make_client(session, "athena", region_name=env.get("REGION", ""))
    response = athena.start_query_execution(
        QueryString=query_sql,
        WorkGroup=env.get("ENVIRONMENT", "dev"),
    )
    Path("response.json").write_text(
        json.dumps(response, default=str), encoding="utf-8"
    )
    print_message(json.dumps(response, default=str))
    query_id = str(response.get("QueryExecutionId", ""))
    if not query_id:
        raise RuntimeError("[Dados] Falha ao iniciar query no Athena")
    print_message(f"QueryExecutionId: {query_id}")
    return query_id


def wait_for_athena_query(
    table_name: str, query_id: str, env: dict[str, str], verification: str
) -> tuple[bool, bool]:
    """
    Aguarda a conclusão de uma query no Athena e verifica os resultados.

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

    Returns:
        tuple[bool, bool]: Um tupla indicando o status da query e da verificação.

    """
    session = make_session(env)
    athena = make_client(session, "athena", region_name=env.get("REGION", ""))
    s3 = make_client(session, "s3", region_name=env.get("REGION", ""))

    while True:
        body = athena.get_query_execution(QueryExecutionId=query_id)
        body_text = json.dumps(body, default=str)
        Path("execution.json").write_text(body_text, encoding="utf-8")
        print_message(body_text)
        status = (
            body.get("QueryExecution", {}).get("Status", {}).get("State", "").upper()
        )

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

        if status in {"QUEUED", "RUNNING"}:
            print_message(
                f"[Dados] athena - {table_name} ainda em execução, tentando novamente em "
                f"{WAIT_FOR_ATHENA_QUERY_TIME}s..."
            )
            time.sleep(WAIT_FOR_ATHENA_QUERY_TIME)
            continue

        output_location = (
            body.get("QueryExecution", {})
            .get("ResultConfiguration", {})
            .get("OutputLocation", "")
        )
        print_message(f"QueryExecutionId: {output_location}")
        if not output_location:
            return False, False

        csv_file = Path(f"{table_name}.csv")
        time.sleep(5)
        parsed = urlparse(output_location)
        bucket = parsed.netloc
        key = parsed.path.lstrip("/")
        s3.download_file(bucket, key, str(csv_file))
        line_count = len(
            csv_file.read_text(encoding="utf-8", errors="ignore").splitlines()
        )
        print_message(f"[Dados] athena - {table_name} qtd: {line_count}")
        if line_count >= 2:
            print_message(
                f"[Dados] athena - {verification} encontrado, invalidando CI!"
            )
            return False, True

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


def start_and_wait_athena_queries(
    pairs: list[tuple[str, str]], env: dict[str, str], verification: str
) -> None:
    """
    Aguarda a conclusão de múltiplas queries no Athena e verifica os resultados.

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

    Raises:
        RuntimeError: Se ocorrer um erro na execução de alguma query ou se forem encontrados dados duplicados.

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

    print_message(f"[Dados] athena - Aguardando {WAIT_START_ATHENA_QUERY_TIME}...")
    time.sleep(WAIT_START_ATHENA_QUERY_TIME)

    errors: list[str] = []
    duplicates: list[str] = []
    for table_name, query_id in query_ids:
        has_error, has_dup = wait_for_athena_query(
            table_name, query_id, env, verification
        )
        if has_error:
            errors.append(table_name)
        if has_dup:
            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] athena ok!")