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.models.baseoperator import chain
from airflow.models.taskinstance import TaskInstance
from airflow.providers.mysql.hooks.mysql import MySqlHook
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
def get_task_list(ti: TaskInstance):
    db = MySqlHook(mysql_conn_id="autonomous-serverless-instance")
    task_list = get_records_from_db(
        db,
        f"""
        select * from auto_deploy_dedicated_{ENV}.register_default_sku_tasks
        """,
    )
    ti.xcom_push("task_list", task_list)
    return len(task_list) > 0


@task
def notify_task_begin(ti: TaskInstance):
    task_list = ti.xcom_pull(key="task_list")
    for t in task_list:
        if not t["trigger"]:
            continue
        lark.send_card_message(
            lark.build_status_card(
                lark.Color.BLUE,
                f"start registering spu for region {t['logical_region_name']}",
                {
                    "☁️ Cloud": t["provider"],
                    "📍 Region": t["region_name"],
                    "🪪 Name": t["logical_region_name"],
                    "📦 SPUs Collection": t["spus_collection_name"],
                },
                {
                    "SKU Service": "registering",
                },
            ),
            email=t["trigger"],
        )


@task
def notify_task_end(ti: TaskInstance):
    task_list = ti.xcom_pull(key="task_list")

    for t in task_list:
        if not t["trigger"]:
            continue
        lark.send_card_message(
            lark.build_status_card(
                lark.Color.GREEN,
                f"finish registering spu for region {t['logical_region_name']}",
                {
                    "☁️ Cloud": t["provider"],
                    "📍 Region": t["region_name"],
                    "🪪 Name": t["logical_region_name"],
                    "📦 SPUs Collection": t["spus_collection_name"],
                    "🏁 Status": "SUCCESS",
                },
                {
                    "SKU Service": "active",
                },
            ),
            email=t["trigger"],
        )


@task(
    outlets=[
        Dataset(
            f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/default_sku_services"
        )
    ]
)
def update_task_status(ti: TaskInstance):
    task_list = ti.xcom_pull(key="task_list")
    db = MySqlHook(mysql_conn_id="autonomous-serverless-instance")
    with closing(db.get_conn()) as conn, closing(conn.cursor()) as cur:
        for t in task_list:
            cur.execute(
                f"""
                insert into auto_deploy_dedicated_{ENV}.register_default_sku_results 
                    (request_id, provider, region_name, logical_region_name, status, trace_id)
                values (%(request_id)s, %(provider)s, %(region_name)s, %(logical_region_name)s, %(status)s, %(trace_id)s)
                """,
                {
                    **t,
                    "request_id": t["id"],
                    "status": "success",
                },
            )
            cur.execute(
                f"""
                insert into auto_deploy_dedicated_{ENV}.default_sku_services 
                    (provider, region_name, logical_region_name, status, trace_id, requested_at)
                values (%(provider)s, %(region_name)s, %(logical_region_name)s, %(status)s, %(trace_id)s, %(requested_at)s)
                """,
                {
                    **t,
                    "requested_at": t["created_at"],
                    "status": "active",
                },
            )
        conn.commit()


@dag(
    dag_id=f"register_default_sku_{ENV}",
    start_date=pendulum.datetime(2024, 11, 26, tz="UTC"),
    schedule=[
        Dataset(
            f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/register_default_sku_requests"
        )
    ],
    tags=["default-sku", "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 register_default_sku():
    chain(
        get_task_list(),
        notify_task_begin(),
        github_action(
            group_id="register_default_sku",
            repo="tidbcloud/auto-deploy",
            workflow="register-default-skus.yml",
            ref=AUTO_DEPLOY_REPO_BRANCH,
            input_fn=lambda ti: map(
                lambda t: {
                    "env": ENV,
                    "cloud": t["provider"],
                    "region": t["region_name"],
                    "spus-collection-name": t["spus_collection_name"],
                    "sku-collection-name": t["sku_collection_name"],
                    "recreate": "false",
                },
                ti.xcom_pull(key="task_list"),
            ),
        ),
        update_task_status(),
        notify_task_end(),
    )


register_default_sku()