import hashlib
import logging
from contextlib import closing

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,
    write_o11y_config_requests,
    write_o11y_create_requests,
)
from dotenv import dotenv_values
from utils.aws_az_check import find_unsupported_az_ids
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
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 byoc region infomation",
                {
                    "☁️ Cloud": t["provider"],
                    "📍 Region": t["region_name"],
                    "🪪 Name": t["logical_region_name"],
                    "🧩 AZ IDs": t["availability_zones"],
                },
                {
                    "ControlPlane Account": t["m_account"],
                    "DataPlane Account": t["d_accounts"],
                    "O11Y Account": t["o11y_account"],
                },
            ),
            email=t["trigger"],
        )


@task.short_circuit(ignore_downstream_trigger_rules=False)
def get_create_byoc_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}.byoc_region_create_tasks
        """,
    )
    for r in requests:
        r["is_ng"] = True
        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
        unsupported_az_ids = find_unsupported_az_ids(r["availability_zones"])
        if unsupported_az_ids:
            logger.error(
                f"Unsupported AZ IDs found in {r['region_name']}: {unsupported_az_ids}"
            )
            raise AirflowFailException(
                f"Unsupported AZ IDs found: {', '.join(unsupported_az_ids)}"
            )

    # 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(
    outlets=[
        Dataset(
            f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/byoc_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}.byoc_regional_create_requests
                    (provider, region_name, logical_region_name, rds_instance_name, create_rds, `trigger`, m_account, d_accounts, o11y_account, trace_id, availability_zones, advance_options)
                values (%(provider)s, %(region_name)s, %(logical_region_name)s, %(rds_instance_name)s, %(create_rds)s, %(trigger)s, %(m_account)s, %(d_accounts)s, %(o11y_account)s, %(trace_id)s, %(availability_zones)s, %(advance_options)s)
                """,
                r,
            )
        conn.commit()


@task(
    outlets=[
        Dataset(
            f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/byoc_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}.byoc_region_create_results
                    (request_id, trace_id) 
                values (%(id)s, %(trace_id)s)
                """,
                r,
            )
        conn.commit()


@dag(
    dag_id=f"deploy_new_byoc_region_{ENV}",
    start_date=pendulum.datetime(2025, 5, 20, tz="UTC"),
    schedule=[
        Dataset(
            f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/byoc_region_create_requests"
        )
    ],
    tags=["deploy", "dedicated", "region", "byoc", 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_byoc_region():
    chain(
        get_create_byoc_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"],
                }
                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"
            ],
        ),
        [
            diff_regional_create_requests(ENV),
            diff_o11y_create_requests(),
        ],
        [
            write_regional_create_requests(),
            write_o11y_create_requests(),
        ],
        compute_o11y_config_requests(
            f"auto_deploy_dedicated_{ENV}.byoc_regional_create_requests"
        ),
        write_o11y_config_requests(ENV),
        write_task_results(),
    )


deploy_new_byoc_region()