import logging
from contextlib import closing

import pendulum
import requests
import utils.lark as lark
from airflow.datasets import Dataset
from airflow.decorators import dag, task
from airflow.exceptions import AirflowFailException
from airflow.models import Variable
from airflow.models.baseoperator import chain
from airflow.models.taskinstance import TaskInstance
from airflow.providers.mysql.hooks.mysql import MySqlHook
from airflow.sensors.base import PokeReturnValue
from dotenv import dotenv_values

MYSQL_CONN_ID = "autonomous-serverless-instance"
DB_INSTANCE = "auto_deploy_o11y_prod"
AWS_ROLE = "auto-deploy-cli"

GIT_REPO = "tidbcloud/auto-deploy"
GIT_REF = "prod"


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

task_logger = logging.getLogger("airflow.task")


@task.short_circuit
def get_task_list(ti: TaskInstance):
    """
    fetch destroy regional tasks from o11y_destroy_tasks
    """
    db = MySqlHook(mysql_conn_id=MYSQL_CONN_ID)
    tasks = db.get_pandas_df(
        f"select * from {DB_INSTANCE}.o11y_destroy_tasks",
    ).to_dict("records")
    if len(tasks) > 1:
        task_logger.error(f"Too many task: {len(tasks)}")
        raise AirflowFailException("Too many task")
    elif len(tasks) == 1:
        ti.xcom_push(key="task", value=tasks[0])
        return True
    else:
        return False


@task
def get_region_meta(ti: TaskInstance):
    """
    filter task, if task is invalid update task result
    """
    task = ti.xcom_pull(key="task")
    db = MySqlHook(mysql_conn_id=MYSQL_CONN_ID)
    meta = db.get_pandas_df(
        f"""
        select meta_uri from {DB_INSTANCE}.o11y_regions
        where provider = %(provider)s and logical_region_name = %(logical_region_name)s
        """,
        task,
    ).to_dict("records")
    if len(meta) == 0:
        task_logger.error("no meta for %(provider)s", task)
        task["status"] = "invalid"
    else:
        task_logger.info("get meta-uri: %(meta_uri)s", meta[0])
        task["meta_uri"] = meta[0]["meta_uri"]
        task["status"] = None
    ti.xcom_push(key="task", value=task)


@task
def call_github_action(ti: TaskInstance):
    """
    call github action destroy-aws-region-k8s
    """
    token = Variable.get("git_action_bot_token")
    task = ti.xcom_pull(key="task")
    if task["status"] is not None:
        return

    ti.xcom_push(key="github_action_begin", value=pendulum.now().to_iso8601_string())
    r = requests.post(
        f"https://api.github.com/repos/{GIT_REPO}/actions/workflows/destroy-o11y-aws-region-k8s.yml/dispatches",
        json={
            "ref": GIT_REF,
            "inputs": {
                "env": "prod",
                "region": task["logical_region_name"],
                "meta-uri": task["meta_uri"],
                "role-arn": f"arn:aws:iam::{task['o11y_account']}:role/{AWS_ROLE}",
            },
        },
        headers={
            "Authorization": f"Bearer {token}",
            "Accept": "application/vnd.github+json",
            "X-GitHub-Api-Version": "2022-11-28",
        },
    )
    if r.status_code != 204:
        # TODO:
        #    1. alert to oncaller
        #    2. log the error with distributed tracing system
        task["status"] = "abort"
        task_logger.error(f"run workflow response: {r.status_code}:{r.text}")
    else:
        task_logger.error(
            "run workflow success %(provider)s %(logical_region_name)s", task
        )
        ti.xcom_push(
            key="github_action_end",
            value=pendulum.now().add(seconds=10).to_iso8601_string(),
        )

    ti.xcom_push(key="task", value=task)


@task.sensor(poke_interval=30, timeout=3600, mode="reschedule")
def wait_github_action_done(ti: TaskInstance):
    """
    wait all destroy-o11y-aws-region-k8s to be in final state
    """
    in_process_status = [
        "action_required",
        "in_progress",
        "queued",
        "requested",
        "waiting",
        "pending",
    ]
    task = ti.xcom_pull(key="task")
    if task["status"] is not None:
        return PokeReturnValue(is_done=True)

    begin_time = ti.xcom_pull(key="github_action_begin")
    end_time = ti.xcom_pull(key="github_action_end")
    token = Variable.get("git_action_bot_token")
    r = requests.get(
        f"https://api.github.com/repos/{GIT_REPO}/actions/workflows/destroy-o11y-aws-region-k8s.yml/runs",
        params={
            "created": f"{begin_time}..{end_time}",
            "branch": GIT_REF,
        },
        headers={
            "Authorization": f"Bearer {token}",
        },
    )
    if r.status_code != 200:
        task_logger.error(f"get workflow response: {r.status_code}:{r.text}")
        raise AirflowFailException(r.text)
    data = r.json()
    runs = data["workflow_runs"]
    # print data len for debug
    print(f"got {len(runs)} runs from github")
    if len(runs) == 0:
        return PokeReturnValue(is_done=False)
    for run in runs:
        if run["status"] in in_process_status:
            task_logger.info(f"get workflow status: {run['status']}")
            return PokeReturnValue(is_done=False)
        elif run["status"] == "completed":
            task["status"] = "success"
        else:
            task_logger.warning(f"get workflow status: {run['status']}")
            task["status"] = "failure"
    ti.xcom_push(key="task", value=task)
    return PokeReturnValue(is_done=True)


@task(
    outlets=[
        Dataset(f"mysql://{MYSQL_CONN_ID}/{DB_INSTANCE}/o11y_destroy_results"),
        Dataset(f"mysql://{MYSQL_CONN_ID}/{DB_INSTANCE}/o11y_regions"),
    ]
)
def write_task_results(ti: TaskInstance = None):
    task = ti.xcom_pull(key="task")
    db = MySqlHook(mysql_conn_id=MYSQL_CONN_ID)
    with closing(db.get_conn()) as conn, closing(conn.cursor()) as cur:
        cur.execute(
            f"""
            insert into {DB_INSTANCE}.o11y_destroy_results 
                (request_id, status, provider, region_name, trace_id) 
            values (%(id)s, %(status)s, %(provider)s, %(region_name)s, %(trace_id)s)
            """,
            task,
        )
        if task["status"] == "success":
            task["status"] = "inactive"
            cur.execute(
                f"""
                delete from {DB_INSTANCE}.o11y_regions
                where provider = %(provider)s and region_name = %(region_name)s
                """,
                task,
            )
        conn.commit()


@dag(
    start_date=pendulum.datetime(2024, 10, 25, tz="UTC"),
    schedule=[Dataset(f"mysql://{MYSQL_CONN_ID}/{DB_INSTANCE}/o11y_destroy_requests")],
    tags=["destroy", "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 destroy_o11y_region():
    chain(
        get_task_list(),
        get_region_meta(),
        call_github_action(),
        wait_github_action_done(),
        write_task_results(),
    )


destroy_o11y_region()