"""
Modulo principal para o CI/CD do Data Lake, incluindo empacotamento, implantação e publicação de SQL para Redshift.

Este módulo fornece funções para empacotar e publicar SQLs do projeto, materializando objetos no banco de dados,
sincronizar arquivos SQL com o S3 e publicar SQL para Redshift. 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 processos de empacotamento, implantação e
publicação sejam realizados de forma confiável e consistente.

"""

from __future__ import annotations

import argparse
import json
import tempfile
from pathlib import Path

from colorama import Fore

from common import print_message
from src.datalake.commons.aws_utils import (
    configure_aws_access,
    make_client,
    make_session,
    sync_directory_to_s3,
)
from src.datalake.commons.data_env import (
    export_to_process,
    load_runtime_env,
    set_all_env,
)
from src.datalake.commons.lambda_utils import invoke_lambda

FILE_EXTENSIONS_SELECTION_SQL = "*.sql"


def _replace_sql_date_token() -> None:
    sql_dir = Path("sql")
    if not sql_dir.exists():
        return
    for sql_file in sql_dir.rglob(FILE_EXTENSIONS_SELECTION_SQL):
        text = sql_file.read_text(encoding="utf-8")
        text = text.replace("@@SQL_DATE@@", "${flyway:timestamp}")
        sql_file.write_text(text, encoding="utf-8")


def deploy_sql_to_s3(env: dict[str, str], phase: str = "all") -> None:
    sql_dir = Path("sql")
    if not sql_dir.exists():
        return

    print_message(f"{Fore.YELLOW}[Dados] Efetuando upload dos SQLs para S3")

    _replace_sql_date_token()

    sql_path = Path("sql")
    if sql_path.exists() and sql_path.is_dir():
        for sql_dir in sql_path.iterdir():
            print_message(
                f"{Fore.YELLOW}[Dados] Verificando diretório para upload dos SQLs para S3 - {sql_dir}"
            )
            if sql_dir.is_dir() and sql_dir.name not in [
                "post_consolidate",
                "validation",
            ]:
                print_message(
                    f"{Fore.YELLOW}[Dados] Upload dos SQLs para S3 - {sql_dir}"
                )
                sync_directory_to_s3(
                    str(sql_dir),
                    env["WORK_BUCKET_NAME"],
                    f"{env['SOURCE_SYSTEM']}/{env.get('TEAM')}/sql/{phase}/db/migration/{sql_dir.name}/",
                    env,
                    delete=True,
                    exclude=["*"],
                    include=[FILE_EXTENSIONS_SELECTION_SQL],
                )
            else:
                print_message(
                    f"{Fore.YELLOW}[Dados] Ignorando upload dos SQLs para S3 - {sql_dir}"
                )

    deploy_subscriptions_to_s3(env, phase)


def deploy_subscriptions_to_s3(env: dict[str, str], phase: str) -> None:
    session = make_session(env)
    s3 = make_client(session, "s3", region_name=env.get("REGION", ""))
    bucket = env["WORK_BUCKET_NAME"]

    subscriptions = env.get("REDSHIFT_SUBSCRIPTIONS", "")
    print_message(
        f"{Fore.YELLOW}[Dados] Dependentes de SQL para Redshift: '{subscriptions}'"
    )
    if subscriptions:
        print_message(f"{Fore.YELLOW}[Dados] Ativando dependentes de SQL para Redshift")

        filename = f"{env['SOURCE_SYSTEM']}_{env.get('TEAM')}_sql_{phase}.config"

        with tempfile.NamedTemporaryFile(
            mode="w+", encoding="utf-8", delete=True
        ) as tmp_file:
            tmp_file.write(
                f"s3://{bucket}/{env['SOURCE_SYSTEM']}/{env.get('TEAM')}/sql/{phase}/\n",
            )
            for sub in [x.strip() for x in subscriptions.split(",") if x.strip()]:
                dep_key = f"{sub}/sql/{phase}/dependents/{filename}"
                s3.upload_file(str(tmp_file.name), bucket, dep_key)
                print_message(f"[Dados] Dependente de SQL enviado para S3 - {dep_key}")


def publish_redshift(
    type_: str | None = None,
    schemas_list: list[str] | None = None,
    phase: str | None = None,
    extra: str | None = None,
    publish: bool | None = True,
) -> None:
    env = load_runtime_env(
        apply_set_all=True, configure_aws=False, ignore_project_type=True
    )
    env = configure_aws_access(env)
    env = set_all_env(env)
    export_to_process(env)

    print_message(
        f"{Fore.YELLOW}[Dados] Iniciando publicação de SQL para Redshift - "
        f"type: {type_} - schemas: {schemas_list} - phase: {phase} - extra: {extra} - publish: {publish}"
    )

    if not type_:
        type_ = env.get("DATA_TYPE", "invalid")

    if not schemas_list:
        schemas_list = [f"{env.get('DATA_TYPE', '')}_{env.get('SOURCE_SYSTEM', '')}"]

    if not phase:
        phase = env.get("PHASE", "invalid")

    db_name = f"{env.get('ENVIRONMENT')}_{phase}_{env['SOURCE_SYSTEM']}"
    s3_name_name = f"{phase.upper()}_BUCKET_NAME"

    if phase == "stage":
        env["CUSTOM_GRANTS"] = "true"
        env["CUSTOM_S3_SCHEMAS"] = "true"

    if type_ == "tenant":
        db_name = f"{db_name}_tenant"
        s3_name_name = f"{s3_name_name}_TENANT"

    s3_name = env.get(s3_name_name, "")
    if extra == "nodb":
        db_name = ""
    do_publish = "true" if publish else "false"

    env["DB_NAME"] = db_name
    env["S3_NAME"] = s3_name
    env["DO_PUBLISH"] = do_publish

    deploy_sql_to_s3(env, phase=phase or "all")

    print_message(
        f"{Fore.YELLOW}[Dados] Publicando SQL para Redshift - "
        f"type: '{type_}' - schemas: '{schemas_list}' - phase: '{phase}' - extra: {extra} - publish: {publish}"
    )

    lambda_redshift_id = (
        f"arn:aws:lambda:us-east-1:{env.get('AWS_ACCOUNT_ID')}:"
        f"function:{env.get('ENVIRONMENT')}-setup-db-redshift-migrate-deploy-sql"
    )

    schema_list = schemas_list
    if not schema_list:
        schema_list = [""]

    session = make_session(env)
    s3 = make_client(session, "s3", region_name=env.get("REGION", ""))
    for idx, schema in enumerate(schema_list):
        payload = {
            "redo": "false",
            "repair": "false",
            "debug": "true",
            "noSSM": "false",
            "ignoreIgnoredMigrations": "true",
            "sourceSystem": env["SOURCE_SYSTEM"],
            "dataType": type_,
            "schemas": schema,
            "dbName": db_name,
            "stage": phase,
            "s3Name": s3_name,
            "s3SQLEnabled": True,
            "controlPerSchema": True,
            "publish": do_publish if idx == len(schema_list) - 1 else "false",
            "customGrants": env.get("CUSTOM_GRANTS", ""),
            "customS3Schemas": env.get("CUSTOM_S3_SCHEMAS", ""),
        }
        payload_file = Path("setupRedshiftPayload.json")
        payload_file.write_text(
            json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"
        )

        payload_key = (
            f"{env['SOURCE_SYSTEM']}/{env.get('TEAM')}/sql/{phase}/"
            f"setupRedshiftPayload/{type_}_{schema}.json"
        )
        s3.upload_file(str(payload_file), env["WORK_BUCKET_NAME"], payload_key)

        deploy_ignore_time_out = (
            env.get("SOPS_DATACORE_REDSHIFT_DEPLOY_IGNORE_TIME_OUT", "false").lower()
            == "true"
        )
        invoke_lambda(
            lambda_redshift_id,
            f"file://{payload_file}",
            env,
            ignore_timeout=deploy_ignore_time_out,
            description=f"Deploy Redshift - type: '{type_}' - schema: '{schema}' - phase: '{phase}'",
        )


def main() -> int:
    parser = argparse.ArgumentParser(description="Data CI main")

    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument(
        "-d", "--deploy", action="store_true", help="Deploy Redshift changes"
    )

    parser.add_argument(
        "-t",
        "--type",
        required=True,
        choices=["commons", "external", "anonymous", "tenant", "all"],
        help=(
            "Execution type (external, anonymous, tenant, or all) - "
            "will be used to set SOURCE_SYSTEM variable in the env"
        ),
    )
    parser.add_argument(
        "-s",
        "--schemas",
        required=False,
        help="Comma-separated list of schemas to include",
    )
    parser.add_argument(
        "-p",
        "--phase",
        required=False,
        choices=["transient", "raw", "stage", "analytics", "consolidate", "ml", "all"],
        help="Execution phase (transient, raw, stage, analytics, consolidate, ml, all)",
    )
    parser.add_argument(
        "-e",
        "--extra",
        required=False,
        help="Extra parameters for the deployment",
    )
    args = parser.parse_args()

    try:
        if args.deploy:
            type_ = args.type
            schemas = args.schemas.split(",") if args.schemas else None
            phase = args.phase
            extra = args.extra
            publish_redshift(type_, schemas, phase, extra)
        else:
            parser.print_help()
            return 2
    except Exception as exc:
        print_message(str(exc))
        return 2

    return 0


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