"""
Utilitários para interagir com o AWS Glue, incluindo iniciar jobs, aguardar a conclusão e gerenciar payloads.
"""

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_GLUE_JOB_TIME = 30
WAIT_START_GLUE_JOB_TIME = 5


def _job_name(job_name: str, env: dict[str, str]) -> str:
    """
    Retorna o nome completo do job do Glue, incluindo prefixos específicos do ambiente.

    Args:
        job_name (str): Nome do job do Glue.
        env (dict[str, str]): Dicionário contendo as variáveis de ambiente.

    Returns:
        str: Nome completo do job do Glue.

    """
    if job_name.startswith(f"{env.get('ENVIRONMENT')}_"):
        return job_name

    if env.get("THIRD_PARTY") == "true":
        return f"{env.get('ENVIRONMENT')}_tp_{env.get('PHASE')}_{env.get('TEAM')}_{env['SOURCE_SYSTEM']}_{job_name}"
    return (
        f"{env.get('ENVIRONMENT')}_{env.get('PHASE')}_{env['SOURCE_SYSTEM']}_{job_name}"
    )


def start_glue_job(
    job_name: str, env: dict[str, str], payload_dict: dict[str, str] | None = None
) -> str:
    """
    Inicia a execução de um job no AWS Glue.

    Args:
        job_name (str): Nome do job do Glue.
        env (dict[str, str]): Dicionário contendo as variáveis de ambiente.
        payload_dict (dict[str, str] | None): Dicionário contendo os argumentos do job.

    Returns:
        str: ID da execução do job do Glue.

    """
    to_run = _job_name(job_name, env)
    payload = Path("payload.json")
    if not payload_dict:
        payload_text = payload.read_text(encoding="utf-8") if payload.exists() else "{}"
        payload_dict = json.loads(payload_text) if payload_text.strip() else {}
        if not isinstance(payload_dict, dict):
            payload_dict = {}

    print_message(
        f"[Dados] glue - start-job-run: {to_run}: '{json.dumps(payload_dict, ensure_ascii=False, indent=2)}'"
    )
    session = make_session(env)
    glue = make_client(session, "glue", region_name=env.get("REGION", ""))
    response = glue.start_job_run(
        JobName=to_run,
        Arguments=payload_dict,
    )
    response_text = json.dumps(response, default=str)
    Path("response.json").write_text(response_text, encoding="utf-8")
    print_message(response_text)
    run_id = str(response.get("JobRunId", ""))
    if not run_id:
        raise RuntimeError("[Dados] Falha ao iniciar job do Glue")
    return run_id


def wait_for_glue_job(
    job_name: str, run_id: str, env: dict[str, str], error_list: list[str]
) -> None:
    """
    Aguarda a conclusão de um job no AWS Glue.

    Args:
        job_name (str): Nome do job do Glue.
        run_id (str): ID da execução do job do Glue.
        env (dict[str, str]): Dicionário contendo as variáveis de ambiente.
        error_list (list[str]): Lista para armazenar mensagens de erro.

    Returns:
        None

    """
    to_run = _job_name(job_name, env)
    session = make_session(env)
    glue = make_client(session, "glue", region_name=env.get("REGION", ""))

    while True:
        response = glue.get_job_run(JobName=to_run, RunId=run_id)
        response_text = json.dumps(response, default=str)
        print_message(response_text)
        state = response.get("JobRun", {}).get("JobRunState", "")
        state_upper = state.upper()
        if state_upper in {"FAILED", "TIMEOUT", "ERROR"}:
            error_list.append(f"job={to_run} state={state_upper}")
            return
        if state_upper in {"STARTING", "RUNNING", "STOPPING"}:
            print_message(
                f"[Dados] glue - {to_run} ainda em execução, tentando novamente em {WAIT_FOR_GLUE_JOB_TIME}s..."
            )
            time.sleep(WAIT_FOR_GLUE_JOB_TIME)
            continue

        print_message(f"[Dados] glue - {to_run} ok!")
        return


def start_and_wait_glue_jobs(
    job_names: list[str],
    env: dict[str, str],
    payload_dict: dict[str, str] | None = None,
) -> None:
    """
    Inicia e aguarda a conclusão de múltiplos jobs no AWS Glue.

    Args:
        job_names (list[str]): Lista com os nomes dos jobs do Glue.
        env (dict[str, str]): Dicionário contendo as variáveis de ambiente.
        payload_dict (dict[str, str] | None): Dicionário contendo os argumentos dos jobs.

    Returns:
        None

    """
    runs: list[tuple[str, str]] = []
    for job_name in job_names:
        run_id = start_glue_job(job_name, env, payload_dict)
        runs.append((job_name, run_id))
        print_message(f"[Dados] glue - Aguardando {WAIT_START_GLUE_JOB_TIME}...")
        time.sleep(WAIT_START_GLUE_JOB_TIME)

    error_list = []
    for job_name, run_id in runs:
        wait_for_glue_job(job_name, run_id, env, error_list)

    if error_list:
        raise RuntimeError(f"[Dados] glue - Erros encontrados nos jobs: {error_list}")

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