from contextlib import closing
from os.path import dirname, join, realpath

import jinja2
import pendulum
import utils.lark as lark
import yaml
from airflow.datasets import Dataset
from airflow.decorators import dag, task
from airflow.exceptions import AirflowFailException
from airflow.io.path import ObjectStoragePath
from airflow.models.baseoperator import chain
from airflow.models.taskinstance import TaskInstance
from airflow.providers.mysql.hooks.mysql import MySqlHook
from deploy.dedicated.common import get_o11y_instance_type
from dotenv import dotenv_values
from utils.database import get_records_from_db
from utils.github import github_action
from utils.url import blob_url_to_abfs

config = dotenv_values()
LARK_ALERT_CHAT_ID = config["LARK_ALERT_CHAT_ID"]
META_BUCKET_AWS = config["META_BUCKET_AWS"]
META_STORAGE_ACCOUNT = config["META_STORAGE_ACCOUNT"]
DMS_SNS_TOPIC = config["DMS_SNS_TOPIC"]


@task.short_circuit
def get_task_list(ti: TaskInstance):
    db = MySqlHook(mysql_conn_id="autonomous-serverless-instance")
    tasks = get_records_from_db(
        db,
        """
        select *
        from auto_deploy_o11y_prod.o11y_create_tasks
        where provider in ('aws', 'azure')
        """,
    )
    for t in tasks:
        if t["provider"] == "aws":
            t["meta_uri"] = (
                f"s3://{META_BUCKET_AWS}/o11y-regions/{t['logical_region_name']}/meta.yaml"
            )
        elif t["provider"] == "azure":
            t["meta_uri"] = (
                f"https://{META_STORAGE_ACCOUNT}.blob.core.windows.net/meta/o11y-regions/{t['logical_region_name']}/meta.yaml"
            )
    ti.xcom_push(key="task_list", value=tasks)
    return len(tasks) > 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 creating o11y region {t['logical_region_name']}",
                {
                    "☁️ Cloud": t["provider"],
                    "📍 Region": t["logical_region_name"],
                },
            ),
            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"o11y region {t['logical_region_name']} created",
                {
                    "☁️ Cloud": t["provider"],
                    "📍 Region": t["logical_region_name"],
                    "🏁 Status": "SUCCESS",
                },
            ),
            email=t["trigger"],
        )


def parse_regional_role_name(output_uri: str) -> str:
    file = ObjectStoragePath(
        output_uri,
        conn_id="auto-deploy-cli-aws-o11y-prod",
    )
    if not file.is_file():
        raise AirflowFailException(f"job output uri {output_uri} not exist")
    with file.open() as f:
        x = yaml.safe_load(f)
        return x["regionalServer"]["roleArn"].split("/")[1]


def parse_regional_auzre_client_id(output_uri: str) -> str:
    output_uri = blob_url_to_abfs(output_uri)
    file = ObjectStoragePath(
        output_uri,
        conn_id="auto-deploy-cli-azure-o11y-prod",
    )
    if not file.is_file():
        raise AirflowFailException(f"job output uri {output_uri} not exist")
    with file.open() as f:
        x = yaml.safe_load(f)
        return x["iam"]["regionalServer"]["clientId"]


def parse_private_endpoint(output_uri: str) -> str:
    file = ObjectStoragePath(
        output_uri,
        conn_id="auto-deploy-cli-aws-o11y-prod",
    )
    if not file.is_file():
        return "prometheus-endpoint-a.observability.tidbcloud.com"
    with file.open() as f:
        x = yaml.safe_load(f)
        return x["privateDnsName"]


@task
def parse_action_results(ti: TaskInstance = None):
    task_list = ti.xcom_pull(key="task_list")
    public_ips = {}
    regional_role_name = {}
    private_endpoint = {}
    for t in task_list:
        meta_uri = t["meta_uri"]
        if t["provider"] == "azure":
            meta_uri = blob_url_to_abfs(meta_uri)
        meta = ObjectStoragePath(
            meta_uri,
            conn_id=f"auto-deploy-cli-{t['provider']}-o11y-prod",
        )
        if not meta.is_file():
            print(f"{t['meta_uri']} is not a file")
            continue
        with meta.open() as f:
            x = yaml.safe_load(f)
            public_ips[str(t["id"])] = ",".join(
                x["regions"][t["logical_region_name"]]["k8s"]["public-ips"]
            )
            regional_role_name[str(t["id"])] = (
                parse_regional_role_name(
                    x["regions"][t["logical_region_name"]]["job-outputs"]["basic-iam"][
                        "output"
                    ],
                )
                if t["provider"] == "aws"
                else (
                    parse_regional_auzre_client_id(
                        x["regions"][t["logical_region_name"]]["job-outputs"][
                            "basic-iam"
                        ]["output"],
                    )
                    if t["provider"] == "azure"
                    else ""
                )
            )
            private_endpoint[str(t["id"])] = (
                parse_private_endpoint(
                    x["regions"][t["logical_region_name"]]["job-outputs"][
                        "basic-addon"
                    ]["output"],
                )
                if t["provider"] == "aws"
                else ""
            )
    ti.xcom_push("public_ips", public_ips)
    ti.xcom_push("regional_role_name", regional_role_name)
    ti.xcom_push("private_endpoint", private_endpoint)


@task(
    outlets=[
        Dataset(
            "mysql://autonomous-serverless-instance/auto_deploy_o11y_prod/o11y_create_results"
        ),
        Dataset(
            "mysql://autonomous-serverless-instance/auto_deploy_o11y_prod/o11y_regions"
        ),
    ]
)
def write_task_results(ti: TaskInstance = None):
    task_list = ti.xcom_pull(key="task_list")
    public_ips = ti.xcom_pull(key="public_ips")
    regional_role_name = ti.xcom_pull(key="regional_role_name")
    private_endpoint = ti.xcom_pull(key="private_endpoint")
    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(
                """
                insert into auto_deploy_o11y_prod.o11y_create_results
                    (request_id, status, meta_uri, provider, region_name, logical_region_name, trace_id) 
                values (%(id)s, 'success', %(meta_uri)s, %(provider)s, %(region_name)s, %(logical_region_name)s, %(trace_id)s)
                """,
                t,
            )
            tid = str(t["id"])
            cur.execute(
                """
                insert into auto_deploy_o11y_prod.o11y_regions
                    (provider, region_name, logical_region_name, public_ips, regional_role_name, private_endpoint, meta_uri, trace_id, requested_at)
                values (
                    %(provider)s,
                    %(region_name)s,
                    %(logical_region_name)s,
                    %(public_ips)s,
                    %(regional_role_name)s,
                    %(private_endpoint)s,
                    %(meta_uri)s,
                    %(trace_id)s,
                    %(requested_at)s
                )
                """,
                {
                    **t,
                    "public_ips": public_ips[tid],
                    "regional_role_name": regional_role_name[tid],
                    "private_endpoint": private_endpoint[tid],
                    "requested_at": t["created_at"],
                },
            )
        conn.commit()


@dag(
    start_date=pendulum.datetime(2024, 10, 25, tz="UTC"),
    schedule=[
        Dataset(
            "mysql://autonomous-serverless-instance/auto_deploy_o11y_prod/o11y_create_requests"
        )
    ],
    tags=["deploy", "o11y", "region", "prod"],
    max_active_runs=1,
    default_args={
        "on_failure_callback": lambda context: lark.send_error_message_handler(
            LARK_ALERT_CHAT_ID, context
        )
    },
)
def create_o11y_region():
    template_dir = realpath(join(dirname(__file__), "..", "..", "..", "templates"))
    env = jinja2.Environment(
        loader=jinja2.FileSystemLoader(template_dir),
        block_start_string="[%",
        block_end_string="%]",
        variable_start_string="[[",
        variable_end_string="]]",
    )
    aws_template = env.get_template("advance_options_aws_o11y_prod.yaml")
    azure_template = env.get_template("advance_options_azure_o11y_prod.yaml")

    chain(
        get_task_list(),
        notify_task_begin(),
        github_action(
            group_id="deploy-o11y-aws-region",
            repo="tidbcloud/auto-deploy",
            workflow="deploy-o11y-aws-region.yml",
            ref="prod",
            input_fn=lambda ti: [
                {
                    "env": "prod",
                    "region": t["region_name"],
                    "logic-region-name": t["logical_region_name"],
                    "meta-uri": t["meta_uri"],
                    "force-overwrite-resources": False,
                    "advance-options": aws_template.render(
                        output_bucket=META_BUCKET_AWS,
                        o11y_account=t["o11y_account"],
                        d_accounts=t["d_accounts"].split(","),
                        region=t["region_name"],
                        o11y_instance_type=get_o11y_instance_type(
                            "prod", False, t["region_name"]
                        ),
                        enable_rds="true" if t.get("enable_rds", False) else "false",
                    ),
                }
                for t in ti.xcom_pull(key="task_list")
                if t["provider"] == "aws"
            ],
        ),
        github_action(
            group_id="deploy-o11y-azure-region",
            repo="tidbcloud/auto-deploy",
            workflow="deploy-azure-o11y-region.yml",
            ref="prod",
            input_fn=lambda ti: [
                {
                    "tidb-cloud-instance": "prod",
                    "logic-region-name": t["logical_region_name"],
                    "cloud-region": t["region_name"],
                    "meta-uri": t["meta_uri"],
                    # aws_account:tenant_id:subscription_id:client_id
                    "tenant-id": t["o11y_account"].split(":")[1],
                    "subscription-id": t["o11y_account"].split(":")[2],
                    "client-id": t["o11y_account"].split(":")[3],
                    "force-overwrite-resources": False,
                    "advance-options": azure_template.render(
                        output_storage_account=META_STORAGE_ACCOUNT,
                        dms_sns_topic=DMS_SNS_TOPIC,
                        o11y_account=t["o11y_account"].split(":")[0],
                        region=t["region_name"],
                    ),
                }
                for t in ti.xcom_pull(key="task_list")
                if t["provider"] == "azure"
            ],
        ),
        parse_action_results(),
        write_task_results(),
        notify_task_end(),
    )


create_o11y_region()