import json
import logging
from contextlib import closing
from os.path import join

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 dotenv import dotenv_values
from requests.auth import HTTPDigestAuth
from utils.database import get_records_from_db
from utils.github import github_action
from utils.jenkins import jenkins_job

config = dotenv_values()
ENV = config["ENV"]
TEST_URL = config["TEST_URL"]
FULL_TEST_CASES = config["FULL_TEST_CASES"]
SIMPLE_TEST_CASES = config["SIMPLE_TEST_CASES"]
TEST_REQUEST_URL = config["TEST_REQUEST_URL"]
TEST_REDIRECT_URL = config["TEST_REDIRECT_URL"]
DEDICATED_API_ENDPOINT = config["DEDICATED_API_ENDPOINT"]
TEST_TENANT_EMAIL = config["TEST_TENANT_EMAIL"]
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}.region_test_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 running tests on region {t['logical_region_name']}",
                {
                    "☁️ Cloud": t["provider"],
                    "📍 Region": t["region_name"],
                    "🪪 Name": t["logical_region_name"],
                },
            ),
            email=t["trigger"],
        )


@task
def notify_task_end(ti: TaskInstance):
    builds_list = ti.xcom_pull(key="builds")
    status_map = ti.xcom_pull(key="status")
    task_list = ti.xcom_pull(key="task_list")

    for i, t in enumerate(task_list):
        if not t["trigger"]:
            continue
        status = status_map[builds_list[i]]
        lark.send_card_message(
            lark.build_status_card(
                lark.Color.GREEN if status == "SUCCESS" else lark.Color.RED,
                f"run tests on region {t['logical_region_name']}",
                {
                    "☁️ Cloud": t["provider"],
                    "📍 Region": t["region_name"],
                    "🪪 Name": t["logical_region_name"],
                    "🏁 Status": status,
                },
                {
                    "Report": join(builds_list[i], "allure"),
                    "Test Started At": t["created_at"],
                },
            ),
            email=t["trigger"],
        )


def set_cidr_if_not_exist(
    provider: str,
    region_name: str,
    project_id: str,
    cidr_notion: str,
    public_key: str,
    private_key: str,
):
    r = requests.get(
        DEDICATED_API_ENDPOINT + "/v1beta1/networkContainers",
        auth=HTTPDigestAuth(public_key, private_key),
    )
    if r.status_code != 200:
        logger.error(f"fetch cidr list: {r.status_code}:{r.text}")
        raise AirflowFailException(r.text)
    data = r.json()
    containers = data["networkContainers"]
    region_id = f"{provider}-{region_name}"
    for c in containers:
        if c["regionId"] == region_id:
            return
    r = requests.post(
        DEDICATED_API_ENDPOINT + "/v1beta1/networkContainers",
        auth=HTTPDigestAuth(public_key, private_key),
        json={
            "regionId": region_id,
            "cidrNotion": cidr_notion,
            "labels": {
                "tidb.cloud/project": project_id,
            },
        },
    )
    if r.status_code != 200:
        logger.error(f"create cidr: {r.status_code}:{r.text}")
        raise AirflowFailException(r.text)


@task
def setup_project_cidr(ti: TaskInstance):
    task_list = ti.xcom_pull(key="task_list")
    for t in task_list:
        set_cidr_if_not_exist(
            t["provider"],
            t["region_name"],
            t["project_id"],
            "172.16.0.0/20",
            t["api_public_key"],
            t["api_private_key"],
        )


@task(
    outlets=[
        Dataset(
            f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/regional_services"
        )
    ]
)
def update_regional_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:
        cur.execute(
            f"""
            update auto_deploy_dedicated_{ENV}.regional_services set status = 'active'
            where id in %s
            """,
            ([t["regional_id"] for t in task_list],),
        )
        conn.commit()


@task(
    outlets=[
        Dataset(
            f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/region_test_reports"
        )
    ]
)
def write_reports(ti: TaskInstance):
    builds = ti.xcom_pull(key="builds")
    status = ti.xcom_pull(key="status")
    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 i, t in enumerate(task_list):
            cur.execute(
                f"""
                insert into auto_deploy_dedicated_{ENV}.region_test_reports
                    (request_id, requested_at, regional_id, status, report)
                values (%(id)s, %(requested_at)s, %(regional_id)s, %(status)s, %(report)s)
                """,
                {
                    **t,
                    "requested_at": t["created_at"],
                    "status": status[builds[i]],
                    "report": join(builds[i], "allure"),
                },
            )
        conn.commit()


@dag(
    dag_id=f"run_regional_test_{ENV}",
    start_date=pendulum.datetime(2024, 11, 26, tz="UTC"),
    schedule=(
        Dataset(
            f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/region_test_requests"
        )
        | Dataset(
            f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/o11y_config_results"
        )
    ),
    tags=["test", "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 run_regional_test():
    chain(
        get_task_list(),
        notify_task_begin(),
        github_action(
            group_id="add_test_tenant_to_region",
            repo="tidbcloud/auto-deploy",
            workflow="update-region.yml",
            ref=AUTO_DEPLOY_REPO_BRANCH,
            input_fn=lambda ti: map(
                lambda t: {
                    "env": ENV,
                    "logical-name": t["logical_region_name"],
                    "deploy-type": (
                        "DEPLOY_TYPE_NEXTGEN"
                        if ti.xcom_pull(key="task").get("is_ng")
                        else "DEPLOY_TYPE_DEDICATED"
                    ),
                    "options": json.dumps(
                        {"state": "ACTIVE", "openedTenantIds": [t["tenant_id"]]}
                    ),
                },
                ti.xcom_pull(key="task_list"),
            ),
        ),
        update_regional_status(),
        setup_project_cidr(),
        jenkins_job(
            group_id="run_test_cases",
            jenkins_url="http://cicd.tidbcloud.com",
            username="zhoushuai@pingcap.com",
            password=Variable.get("jenkins_token_zhoushuai"),
            job_name="platform-scene-online",
            input_fn=lambda ti: map(
                lambda t: {
                    "URL": TEST_URL,
                    "REQUESTURL": TEST_REQUEST_URL,
                    "ORGID": t["tenant_id"],
                    "PROJECTID": t["project_id"],
                    "TENANTEMAIL": TEST_TENANT_EMAIL,
                    "REDIRECTURL": TEST_REDIRECT_URL,
                    "RUNUSER": "zhoushuai",
                    "SCOPE": (
                        FULL_TEST_CASES if t["scope"] == "full" else SIMPLE_TEST_CASES
                    ),
                    "Region": t["region_name"],
                    "CPU": "4",
                    "HasTiflash": "N",
                    "SPENDLIMITS": "0",
                    "RequireSecureTransport": "off",
                    "Provider": t["provider"],
                    "IsOpenAPI": "N",
                    "EnvFile": ENV,
                    "DoSceneBillsCheck": "N",
                },
                ti.xcom_pull(key="task_list"),
            ),
        ),
        write_reports(),
        notify_task_end(),
    )


run_regional_test()