from contextlib import closing
from typing import List

import yaml
from airflow.datasets import Dataset
from airflow.decorators import task
from airflow.exceptions import AirflowFailException
from airflow.io.path import ObjectStoragePath
from airflow.models.taskinstance import TaskInstance
from airflow.providers.mysql.hooks.mysql import MySqlHook
from airflow.utils.trigger_rule import TriggerRule
from utils.database import get_records_from_db
from utils.github import github_action
from utils.url import blob_url_to_abfs


def get_events_with_trace_ids(trace_ids: List[str], table: str) -> List[dict]:
    db = MySqlHook(mysql_conn_id="autonomous-serverless-instance")
    return get_records_from_db(
        db,
        f"""
            select * from {table}
            where trace_id in %s
            """,
        (trace_ids,),
    )


@task.short_circuit(
    trigger_rule=TriggerRule.NONE_FAILED,
    ignore_downstream_trigger_rules=False,
)
def compute_o11y_config_requests(regional_requests_table: str, ti: TaskInstance = None):
    requests = ti.xcom_pull(key="requests", default=[])
    if len(requests) == 0:
        return False
    trace_ids = [r["trace_id"] for r in requests]
    regional_requests = get_events_with_trace_ids(trace_ids, regional_requests_table)
    o11y_requests = get_events_with_trace_ids(
        trace_ids, "auto_deploy_o11y_prod.o11y_create_requests"
    )
    dedicated_region_names = set([r["logical_region_name"] for r in regional_requests])
    o11y_region_names = set([r["logical_region_name"] for r in o11y_requests])

    requests = list(
        filter(
            lambda r: r["o11y_region_name"] in o11y_region_names
            or r["logical_region_name"] in dedicated_region_names,
            requests,
        )
    )
    ti.xcom_push(key="o11y_config_requests", value=requests)
    return len(requests) > 0


@task.short_circuit(ignore_downstream_trigger_rules=False)
def diff_o11y_create_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_events_with_trace_ids(
            new_trace_ids, "auto_deploy_o11y_prod.o11y_create_requests"
        )
    ]
    existing_regions = [
        rec["logical_region_name"]
        for rec in get_records_from_db(
            db,
            """
            select logical_region_name from auto_deploy_o11y_prod.o11y_regions
            """,
        )
    ]
    requests = [
        r
        for r in requests
        if r["trace_id"] not in existing_trace_ids
        and r["o11y_region_name"] not in existing_regions
    ]
    ti.xcom_push(key="o11y_requests", value=requests)
    return len(requests) > 0


@task.short_circuit(ignore_downstream_trigger_rules=False)
def diff_regional_create_requests(env: str, 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_events_with_trace_ids(
            new_trace_ids,
            f"auto_deploy_dedicated_{env}.regional_create_requests",
        )
    ] + [
        r["trace_id"]
        for r in get_events_with_trace_ids(
            new_trace_ids,
            f"auto_deploy_dedicated_{env}.byoc_regional_create_requests",
        )
    ]
    existing_regions = [
        rec["logical_region_name"]
        for rec in get_records_from_db(
            db,
            f"""
            select logical_region_name from auto_deploy_dedicated_{env}.regional_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="regional_requests", value=requests)
    return len(requests) > 0


def write_o11y_config_requests(env: str):
    @task(
        task_id="write_o11y_config_requests",
        outlets=[
            Dataset(
                f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{env}/o11y_config_requests"
            )
        ],
    )
    def operator(ti: TaskInstance = None):
        requests = ti.xcom_pull(key="o11y_config_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}.o11y_config_requests
                        (m_account, d_accounts, o11y_account, provider, region_name, logical_region_name, o11y_region_name, `trigger`, trace_id)
                    values (%(m_account)s, %(d_accounts)s, %(o11y_account)s, %(provider)s, %(region_name)s, %(logical_region_name)s, %(o11y_region_name)s, %(trigger)s, %(trace_id)s)
                    """,
                    r,
                )
            conn.commit()

    return operator()


@task(
    outlets=[
        Dataset(
            "mysql://autonomous-serverless-instance/auto_deploy_o11y_prod/o11y_create_requests"
        )
    ]
)
def write_o11y_create_requests(ti: TaskInstance = None):
    requests = ti.xcom_pull(key="o11y_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(
                """
                insert into auto_deploy_o11y_prod.o11y_create_requests
                    (provider, region_name, logical_region_name, `trigger`, o11y_account, d_accounts, trace_id)
                values (%(provider)s, %(region_name)s, %(o11y_region_name)s, %(trigger)s, %(o11y_account)s, %(d_accounts)s, %(trace_id)s)
                """,
                {
                    **r,
                },
            )
        conn.commit()


# this is not a task, but a helper function
def get_public_ips(ti: TaskInstance, request_id: int) -> str:
    public_ips = ti.xcom_pull(key="public_ips")
    return public_ips[str(request_id)]


def add_infra_apiserver_source_ips(env):
    return (
        [
            github_action(
                group_id="add_infra_apiserver_source_ip",
                repo="tidbcloud/app-delivery",
                workflow="add-infra-apiserver-source-ip.yaml",
                input_fn=lambda ti: [
                    {
                        "env": env,
                        "ips": get_public_ips(ti, t["id"]),
                        "comment": f"seed {t['provider']} {t['logical_region_name']}",
                    }
                    for t in ti.xcom_pull(key="task_list")
                    if env == "prod"  # only for prod
                ],
            )
        ]
        if env == "prod"
        else []
    )


def get_meta_object(provider: str, meta_uri: str, env: str) -> ObjectStoragePath:
    if provider == "aws":
        return ObjectStoragePath(meta_uri, conn_id=f"auto-deploy-cli-aws-{env}")
    elif provider == "azure":
        meta_uri = blob_url_to_abfs(meta_uri)
        return ObjectStoragePath(meta_uri, conn_id=f"auto-deploy-cli-azure-{env}")
    else:
        raise AirflowFailException(f"unknown provider: {provider}")


@task
def parse_k8s_public_ips(env: str, ti: TaskInstance = None):
    task_list = ti.xcom_pull(key="task_list")
    public_ips = {}
    for t in task_list:
        meta = get_meta_object(t["provider"], t["meta_uri"], env)
        if not meta.is_file():
            continue
        with meta.open() as f:
            try:
                x = yaml.safe_load(f)
                public_ips[str(t["id"])] = ",".join(
                    x["regions"][t["logical_region_name"]]["k8s"]["public-ips"]
                )
            except Exception as e:
                raise AirflowFailException(f"parse public-ips: {e}")
    ti.xcom_push("public_ips", public_ips)


def get_first_availability_zone(zones: str):
    zoneList = zones.split(",")
    zoneList.sort()
    return zoneList[0]


def get_o11y_instance_type(env: str, is_byoc: bool, region: str):
    no_r6a_regions = [
        "af-south-1",
        "ap-east-1",
        "ap-east-2",
        "ap-northeast-2",
        "ap-northeast-3",
        "ap-southeast-3",
        "ap-southeast-4",
        "ap-southeast-5",
        "ap-southeast-6",
        "ap-southeast-7",
        "mx-central-1",
    ]
    return "r6i.xlarge" if region in no_r6a_regions else "r6a.xlarge"


def get_seed_instance_type(env: str, byoc: bool, region: str):
    return "m6i.large"