import hashlib
import logging
from contextlib import closing
from typing import List

import pendulum
import utils.lark as lark
from airflow.datasets import Dataset
from airflow.decorators import dag, task
from airflow.exceptions import AirflowFailException
from airflow.models.baseoperator import chain
from airflow.models.taskinstance import TaskInstance
from airflow.providers.mysql.hooks.mysql import MySqlHook
from airflow.utils.trigger_rule import TriggerRule
from deploy.dedicated.common import (
    compute_o11y_config_requests,
    diff_o11y_create_requests,
    diff_regional_create_requests,
    get_events_with_trace_ids,
    write_o11y_config_requests,
    write_o11y_create_requests,
)
from dotenv import dotenv_values
from utils.database import get_records_from_db
from utils.github import github_action

config = dotenv_values()
ENV = config["ENV"]
LARK_ALERT_CHAT_ID = config["LARK_ALERT_CHAT_ID"]
AUTO_DEPLOY_REPO_BRANCH = config["AUTO_DEPLOY_REPO_BRANCH"]

logger = logging.getLogger(__name__)


@task.short_circuit(ignore_downstream_trigger_rules=False)
def get_create_region_requests(ti: TaskInstance = None):
    db = MySqlHook(mysql_conn_id="autonomous-serverless-instance")
    requests = get_records_from_db(
        db,
        f"""
        select * from auto_deploy_dedicated_{ENV}.region_create_tasks
        """,
    )
    for r in requests:
        r["is_ng"] = False
        if not r["logical_region_name"]:
            m = hashlib.md5()
            m.update(r["trace_id"].encode())
            suffix = m.hexdigest()[-6:]
            r["logical_region_name"] = r["region_name"] + "-" + suffix
    # concurrency is not allowed at this time
    if len(requests) > 1:
        logger.error(f"Too many requests: {len(requests)}")
        raise AirflowFailException("Too many requests")
    ti.xcom_push(key="requests", value=requests)
    return len(requests) > 0


@task
def notify_task_begin(ti: TaskInstance):
    task_list = ti.xcom_pull(key="requests")
    for t in task_list:
        if not t["trigger"]:
            continue
        lark.send_card_message(
            lark.build_status_card(
                lark.Color.BLUE,
                "new region infomation",
                {
                    "☁️ Cloud": t["provider"],
                    "📍 Region": t["region_name"],
                    "🪪 Name": t["logical_region_name"],
                },
                {
                    "ControlPlane Account": t["m_account"],
                    "DataPlane Account": t["d_accounts"],
                    "O11Y Account": t["o11y_account"],
                    "SPUs Collection Name": t["spus_collection_name"],
                    "SKU Collection Name": t["sku_collection_name"],
                },
            ),
            email=t["trigger"],
        )


def get_register_default_sku_requests(trace_ids: List[str]) -> List[dict]:
    return get_events_with_trace_ids(
        trace_ids, f"auto_deploy_dedicated_{ENV}.register_default_sku_requests"
    )


@task.short_circuit(ignore_downstream_trigger_rules=False)
def diff_register_default_sku_requests(ti: TaskInstance = None):
    db = MySqlHook(mysql_conn_id="autonomous-serverless-instance")
    requests = ti.xcom_pull(key="requests")
    new_trace_ids = [r["trace_id"] for r in requests]
    existing_trace_ids = [
        r["trace_id"] for r in get_register_default_sku_requests(new_trace_ids)
    ]
    existing_regions = [
        rec["logical_region_name"]
        for rec in get_records_from_db(
            db,
            f"""
            select logical_region_name
            from auto_deploy_dedicated_{ENV}.default_sku_services
            """,
        )
    ]
    requests = [
        r
        for r in requests
        if r["trace_id"] not in existing_trace_ids
        and r["logical_region_name"] not in existing_regions
    ]
    ti.xcom_push(key="register_default_sku_requests", value=requests)
    return len(requests) > 0


@task(
    outlets=[
        Dataset(
            f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/regional_create_requests"
        )
    ]
)
def write_regional_create_requests(ti: TaskInstance = None):
    requests = ti.xcom_pull(key="regional_requests")
    db = MySqlHook(mysql_conn_id="autonomous-serverless-instance")
    with closing(db.get_conn()) as conn, closing(conn.cursor()) as cur:
        for r in requests:
            cur.execute(
                f"""
                insert into auto_deploy_dedicated_{ENV}.regional_create_requests
                    (provider, region_name, logical_region_name, `trigger`, m_account, d_accounts, o11y_account, trace_id)
                values (%(provider)s, %(region_name)s, %(logical_region_name)s, %(trigger)s, %(m_account)s, %(d_accounts)s, %(o11y_account)s, %(trace_id)s)
                """,
                r,
            )
        conn.commit()


@task(
    outlets=[
        Dataset(
            f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/register_default_sku_requests"
        )
    ]
)
def write_register_default_sku_requests(ti: TaskInstance = None):
    requests = ti.xcom_pull(key="register_default_sku_requests")
    db = MySqlHook(mysql_conn_id="autonomous-serverless-instance")
    with closing(db.get_conn()) as conn, closing(conn.cursor()) as cur:
        for r in requests:
            cur.execute(
                f"""
                insert into auto_deploy_dedicated_{ENV}.register_default_sku_requests
                    (
                        provider, region_name, logical_region_name, `trigger`, spus_collection_name, sku_collection_name, trace_id
                    )
                values
                    (
                        %(provider)s, %(region_name)s, %(logical_region_name)s, %(trigger)s, %(spus_collection_name)s, %(sku_collection_name)s, %(trace_id)s
                    )
                """,
                r,
            )
        conn.commit()


def write_security_record(cur, table: str, record: dict, account: str):
    account_md5 = hashlib.md5(account.encode()).hexdigest()
    cur.execute(
        f"""
        replace into auto_deploy_security.{table}
            (provider, account_md5, region_name, trace_id, requested_at) 
        values (%(provider)s, %(account_md5)s, %(region_name)s, %(trace_id)s, %(requested_at)s)
        """,
        {
            **record,
            "account_md5": account_md5,
            "requested_at": record["created_at"],
        },
    )


@task
def write_security_components(ti: TaskInstance = None):
    requests = ti.xcom_pull(key="requests")
    db = MySqlHook(mysql_conn_id="autonomous-serverless-instance")
    with closing(db.get_conn()) as conn, closing(conn.cursor()) as cur:
        for r in requests:
            for tb in ["security_hub_services"]:
                write_security_record(cur, tb, r, r["m_account"])
                write_security_record(cur, tb, r, r["o11y_account"])
                for account in r["d_accounts"].split(","):
                    account = account.strip()
                    if account == "":
                        continue
                    write_security_record(cur, tb, r, account)
        conn.commit()


@task(
    outlets=[
        Dataset(
            f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/region_create_results"
        )
    ],
    trigger_rule=TriggerRule.NONE_FAILED,
)
def write_task_results(ti: TaskInstance = None):
    requests = ti.xcom_pull(key="requests")
    db = MySqlHook(mysql_conn_id="autonomous-serverless-instance")
    with closing(db.get_conn()) as conn, closing(conn.cursor()) as cur:
        for r in requests:
            cur.execute(
                f"""
                insert into auto_deploy_dedicated_{ENV}.region_create_results
                    (request_id, trace_id) 
                values (%(id)s, %(trace_id)s)
                """,
                r,
            )
        conn.commit()


@dag(
    dag_id=f"deploy_new_region_{ENV}",
    start_date=pendulum.datetime(2024, 10, 25, tz="UTC"),
    schedule=[
        Dataset(
            f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/region_create_requests"
        )
    ],
    tags=["deploy", "dedicated", "region", ENV],
    max_active_runs=1,
    default_args={
        "on_failure_callback": lambda context: lark.send_error_message_handler(
            LARK_ALERT_CHAT_ID, context
        )
    },
)
def deploy_new_region():
    security_components = (
        [
            github_action(
                group_id="enable_security_hub",
                repo="tidbcloud/auto-deploy",
                workflow="enable-aws-security-hub.yaml",
                ref=AUTO_DEPLOY_REPO_BRANCH,
                input_fn=lambda ti: [
                    {
                        "region": r["region_name"],
                        "m-account": r["m_account"],
                        "d-accounts": r["d_accounts"],
                        "enable": "enable",
                    }
                    for r in ti.xcom_pull(key="requests")
                    if r["provider"] == "aws"
                ],
            ),
            write_security_components(),
        ]
        if ENV == "prod"
        else []
    )
    chain(
        get_create_region_requests(),
        notify_task_begin(),
        github_action(
            group_id="enable_aws_dedicated_region",
            repo="tidbcloud/auto-deploy",
            workflow="enable-aws-region.yaml",
            ref=AUTO_DEPLOY_REPO_BRANCH,
            input_fn=lambda ti: [
                {
                    "region": r["region_name"],
                    "m-account": r["m_account"],
                    "d-accounts": r["d_accounts"],
                }
                for r in ti.xcom_pull(key="requests")
                if r["provider"] == "aws"
            ],
        ),
        github_action(
            group_id="enable_aws_o11y_region",
            repo="tidbcloud/auto-deploy",
            workflow="enable-aws-region.yaml",
            ref=AUTO_DEPLOY_REPO_BRANCH,
            input_fn=lambda ti: [
                {
                    "region": r["region_name"],
                    "m-account": r["o11y_account"],
                }
                for r in ti.xcom_pull(key="requests")
                if r["provider"] == "aws"
            ],
        ),
        *security_components,
        [
            diff_regional_create_requests(ENV),
            diff_o11y_create_requests(),
            diff_register_default_sku_requests(),
        ],
        [
            write_regional_create_requests(),
            write_o11y_create_requests(),
            write_register_default_sku_requests(),
        ],
        compute_o11y_config_requests(
            f"auto_deploy_dedicated_{ENV}.regional_create_requests"
        ),
        write_o11y_config_requests(ENV),
        write_task_results(),
    )


deploy_new_region()