import base64
import json
import os
import re
from datetime import datetime, timedelta, timezone
from uuid import uuid4

import aiohttp
import boto3
import tuotuo
from sinks import FullActivityClient, jira_alerts
from tuotuo.activities import Activity, ActivityClient, ActivityType, Sevirity
from tuotuo.logs import LogEntity
from tuotuo.metrics import QueryParam, QueryRange

ACTIVITY_BUCKET = os.environ["ACTIVITY_BUCKET"]
PATH = "tidb_oom_detector"
CE_ENDPOINT = os.environ["CE_ENDPOINT"]
OOM_ANALYSIS_TASK = os.environ["OOM_ANALYSIS_TASK"]
ACTIVITY_HUB_ENDPOINT = os.environ["ACTIVITY_HUB_ENDPOINT"]
SOURCE_KEY = os.environ["SOURCE_KEY"]
SOURCE = os.environ["SOURCE"]
JIRA_ENDPOINT = os.environ.get("JIRA_ENDPOINT")
JIRA_EMAIL = os.environ.get("JIRA_EMAIL")
JIRA_API_TOKEN = os.environ.get("JIRA_API_TOKEN")


async def write_activity_to_hub(activity: Activity):
    async with FullActivityClient(
        source=SOURCE, source_key=SOURCE_KEY, endpoint=ACTIVITY_HUB_ENDPOINT
    ) as client:
        return await client.write_activity(activity)


async def find_oom_events(cluster_id, start, end):
    async with ActivityClient(
        source=SOURCE, source_key=SOURCE_KEY, endpoint=ACTIVITY_HUB_ENDPOINT
    ) as client:
        print(
            f"try to find latest oom event from activity hub, cluster_id: {cluster_id}, start: {start}, end: {end}"
        )
        activities = []
        try:
            activities = await client.search_activities(
                target_id=cluster_id,
                start=start,
                end=end,
            )
        except Exception as e:
            print(f"find oom event from activity hub failed, error: {e}")
        for activity in activities[:]:
            if activity.activity_name.lower() != "tidb oom":
                activities.remove(activity)
                continue
        print(f"find {len(activities)} oom events from activity hub, cluster_id: {cluster_id}")
        return activities


def alert_to_jira(activity: Activity, clinic_org_id: str):
    if not all((JIRA_ENDPOINT, JIRA_EMAIL, JIRA_API_TOKEN)):
        print("[Warning][Jira alert skipped] Jira Cloud credentials are not fully configured")
        return

    alerts = []
    occurrence_time = activity.calibration_time
    content = f'StartsAt: {datetime.fromtimestamp(occurrence_time, tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S")} UTC\n'
    content += "Labels:\n"
    content += f"- cluster_id = {activity.target_id}\n"
    content += f"- instance = {activity.payload['instance']}\n"
    content += "Hint:\n"
    content += f"- Please check https://clinic.pingcap.com/portal/#/orgs/{clinic_org_id}/clusters/{activity.target_id}/report/4/{activity.activity_id} to find more details.\n"
    alert = {
        "project": {"key": "O11Y"},
        "summary": "[PROD] TiDB OOM",
        "description": content,
        "issuetype": {"name": "Task"},
        "priority": {"name": "Major"},
        "components": [{"name": "tidb"}],
        # "customfield_10615": activity.application_id,  # tenant_id
        # "customfield_10616": activity.target_id,  # cluster_id
        "customfield_10157": activity.application_id,  # jira cloud tenant_id
        "customfield_10149": activity.target_id,  # jira cloud cluster_id
    }
    alerts.append(alert)
    try:
        jira_res = jira_alerts(
            jira_url=JIRA_ENDPOINT, email=JIRA_EMAIL, token=JIRA_API_TOKEN, alerts=alerts
        )
    except Exception as error:
        print(f"[Warning][Jira alert failed] {type(error).__name__}")
        return
    print(jira_res)


async def run_oom_analysis(
    activity: Activity,
    clinic_org_id: str,
    provider: str,
    region_name: str,
):
    args = {
        "tags": {"source": "oom_detector"},
        "args": {
            "activity_id": activity.activity_id,
            "cluster_id": activity.target_id,
            "instance": activity.payload["instance"],
            "occurrence_ts": activity.calibration_time,
            "tenant_id": activity.application_id,
            "project_id": activity.concept_id,
            "clinic_org_id": clinic_org_id,
            "provider": provider,
            "region_name": region_name,
        },
    }
    url = f"{CE_ENDPOINT}/tasks/{OOM_ANALYSIS_TASK}/async"
    data = json.dumps(args)
    async with aiohttp.ClientSession(trust_env=True) as session:
        res = await session.post(url, data=data)
        return await res.text()


def dict_to_s3(dict, bucket, key):
    s3 = boto3.resource("s3")
    res = s3.Bucket(bucket).put_object(Key=key, Body=json.dumps(dict))
    return res


def generate_activity(cluster_id, instance, entity, tenant_id, project_id):
    activity_id = str(uuid4())
    occurrence_time = int(entity.time.timestamp())
    detect_time = int(datetime.now().timestamp())
    summary = f"{instance} of cluster {cluster_id}: an OOM occurred."

    if entity.node != "fake":
        details = {
            "cluster_id": cluster_id,
            "instance": instance,
            "tenant_id": tenant_id,
            "project_id": project_id,
            "occurrence_time": occurrence_time,
            "detect_time": detect_time,
            "evidence": {
                "source": "k8s eventrouter",
                "content": [entity.content],
            },
        }
    else:
        details = {
            "cluster_id": cluster_id,
            "instance": instance,
            "tenant_id": tenant_id,
            "project_id": project_id,
            "occurrence_time": occurrence_time,
            "detect_time": detect_time,
            "evidence": {
                "source": "Resource Overview Grafana Dashboard",
                "content": [entity.content],
            },
        }

    res = dict_to_s3(details, ACTIVITY_BUCKET, f"{PATH}/{activity_id}.json")
    print(f"write to s3: {res}")

    payload = {
        "summary": summary,
        "cluster_id": cluster_id,
        "instance": instance,
        "occurrence_time": occurrence_time,
        "detect_time": detect_time,
        "details": {
            "type": "s3",
            "bucket": ACTIVITY_BUCKET,
            "key": f"{PATH}/{activity_id}.json",
        },
    }
    oom_event = Activity(
        activity_type=ActivityType.EVENT,
        activity_name="TiDB OOM",
        target_id=cluster_id,
        calibration_time=occurrence_time,
        source=SOURCE,
        source_key=SOURCE_KEY,
        catalog="OOM",
        severity=Sevirity.WARNING,
        activity_id=activity_id,
        application_id=tenant_id,
        concept_id=project_id,
        report_time=detect_time,
        author="oom-detector",
        context_id=activity_id,
        # context=[],
        payload=payload,
        occurrence_time=occurrence_time,
        detect_time=detect_time,
    )
    return oom_event


async def find_oom_evicted(
    cluster_id: str,
    start: datetime,
    end: datetime,
    tenant_id: str = None,
    project_id: str = None,
    region_name: str = None,
    clinic_org_id: str = None,
    provider: str = None,
    shoot_name: str = None,
    dry_run: bool = False,
):
    # get gcp key
    try:
        s3 = boto3.client("s3")
        key = s3.get_object(Bucket=ACTIVITY_BUCKET, Key="secrets/gcp_secret")["Body"]
        os.environ["GCP_PRIVATE_KEY"] = bytes.decode(base64.b64decode(key.read()))
    except Exception as e:
        print(e)

    args = {"region_name": region_name}
    client = tuotuo.client("logs", **args)
    # client = tuotuo.client("logs")
    oom_events = []
    start_ts = int(start.timestamp())
    end_ts = int(end.timestamp())
    component = "kube-eventrouter" if len(shoot_name) > 0 else "event-exporter"
    async for entity in client.read_single_cluster(
        start_ts,
        end_ts,
        cluster_id=cluster_id,
        project_id=project_id,
        tenant_id=tenant_id,
        region_name=region_name,
        components=[component],
        provider=provider,
    ):
        if "The node was low on resource: memory." in entity.content:
            # if 'Got sys oom event' in entity.content:
            if entity.content.startswith("{"):
                content_dict = json.loads(entity.content)
                if component == "kube-eventrouter":
                    instance = content_dict["event"]["metadata"]["name"].split(".")[0]
                    cluster = content_dict["event"]["metadata"]["namespace"][4:]
                    if (
                        cluster == cluster_id
                        and "tidb"
                        in content_dict["event"]["metadata"]["annotations"]["offending_containers"]
                    ):
                        if content_dict["verb"] == "ADDED":
                            oom_event = generate_activity(
                                cluster_id, instance, entity, tenant_id, project_id
                            )
                            oom_events.append(oom_event)

                        if content_dict["verb"] == "DELETED":
                            continue
                        if (
                            content_dict["verb"] == "UPDATED"
                            and content_dict["event"] == content_dict["old_event"]
                        ):
                            continue
                else:
                    instance = content_dict["metadata"]["name"].split(".")[0]
                    cluster = content_dict["metadata"]["namespace"][4:]
                    if (
                        cluster == cluster_id
                        and "tidb"
                        in content_dict["metadata"]["annotations"]["offending_containers"]
                    ):
                        oom_event = generate_activity(
                            cluster_id, instance, entity, tenant_id, project_id
                        )
                        oom_events.append(oom_event)
            else:
                instance = re.search(r"target=pod/db-tidb-(\d{0,4})", entity.content)
                if instance is None:
                    continue
                else:
                    instance = instance.group(0)[11:]  # cut 'target=pod/'
                    oom_event = generate_activity(
                        cluster_id, instance, entity, tenant_id, project_id
                    )
                    oom_events.append(oom_event)

    # add a new condition to findout the oom event
    # two conditions are mutually exclusive
    kube_pod_info_promql = QueryParam(
        time_arg=float(end_ts),
        query_template='kube_pod_info{namespace=~".*$tidb_cluster", pod=~"db-tidb-([0-9]{1,4})"}',
        cluster_id=cluster_id,
        region_name=region_name,
        clinic_org_id=clinic_org_id,
        provider=provider,
        query_name="kube_pod_info",
        query_labels=["pod"],
    )
    metrics_client = tuotuo.client("metrics")
    async with metrics_client:
        pod_info_res = await metrics_client.query(kube_pod_info_promql)
        if pod_info_res.exception is None and len(pod_info_res) > 0:
            pod_info = {
                "cluster": pod_info_res[0]["metric"]["cluster"],
                "namespace": pod_info_res[0]["metric"]["namespace"],
            }
            query_range = QueryRange(start, end, "15s")
            last_terminated_reason_promql = QueryParam(
                time_arg=query_range,
                cluster_id=cluster_id,
                query_template='sum(kube_pod_container_status_last_terminated_reason{cluster="$cluster", namespace="$namespace", reason="OOMKilled", container="tidb"}) by (pod)',
                cluster=pod_info["cluster"],
                namespace=pod_info["namespace"],
                query_name="terminated_reason",
                query_labels=["pod"],
            )
            last_terminated_reason_res = await metrics_client.query(last_terminated_reason_promql)
            if last_terminated_reason_res.exception is None:
                last_terminated_reason = last_terminated_reason_res.frame
            else:
                last_terminated_reason = None
            restart_count_promql = QueryParam(
                time_arg=query_range,
                cluster_id=cluster_id,
                query_template='increase(kube_pod_container_status_restarts_total{cluster="$cluster", namespace="$namespace", container="tidb"}[15s])',
                cluster=pod_info["cluster"],
                namespace=pod_info["namespace"],
                query_name="restart_count",
                query_labels=["pod"],
            )
            restart_count_res = await metrics_client.query(restart_count_promql)
            if restart_count_res.exception is None:
                restart_count = restart_count_res.frame
            else:
                restart_count = None
            if last_terminated_reason is not None and restart_count is not None:
                oom_reboot = last_terminated_reason.mul(restart_count)
                for col_name, col in oom_reboot.items():
                    for index, item in col.items():
                        if (
                            item > 0.5
                        ):  # if the pod is oom killed and rebooted, the value should be 1.0
                            fake_entity = LogEntity(
                                index,
                                cluster_id,
                                "tidb",
                                col_name,
                                "fake",
                                'The conclusion comes from 2 metrics which are similar to the first 2 panels in Resource Overview Dashboard:\nsum(kube_pod_container_status_last_terminated_reason{cluster="$cluster", namespace="$namespace", reason="OOMKilled", container="tidb"}) by (pod)\nincrease(kube_pod_container_status_restarts_total{cluster="$cluster", namespace="$namespace", container="tidb"}[15s])',
                            )
                            oom_event = generate_activity(
                                cluster_id, col_name, fake_entity, tenant_id, project_id
                            )
                            oom_events.append(oom_event)

    if len(oom_events) > 0:
        already_exist_activities = await find_oom_events(cluster_id, start_ts, end_ts)
        for activity in oom_events:
            already_exist = False
            for exist_activity in already_exist_activities:
                if (
                    activity.payload["instance"] == exist_activity.payload["instance"]
                    and activity.payload["occurrence_time"]
                    == exist_activity.payload["occurrence_time"]
                ):
                    already_exist = True
                    break
            if already_exist:
                print(f"[Info][Already exist] {cluster_id} OOM activity: {activity}")
                oom_events.remove(activity)

    if len(oom_events) > 0 and not dry_run:
        for activity in oom_events:
            res = await write_activity_to_hub(activity)
            alert_to_jira(activity, clinic_org_id)
            analyze_task = await run_oom_analysis(activity, clinic_org_id, provider, region_name)
            print(analyze_task)
            if res.get("status") == "ok":
                print(f"[Info][Write event Successful] {cluster_id} OOM activity: {activity}")
            else:
                print(f"[Error][Write event Failed] {cluster_id} OOM activity: {activity}")
    return oom_events


if __name__ == "__main__":
    import asyncio
    import time

    """
    input: cluster_id, start, end, tenant_id, project_id
    output: oom_events
    excute: write details to s3 and write oom_events to activityhub
    non-goal: alert, invoke consumer
    non-goals are implemented in agent
    """
    cluster_id = "1379661944629574082"
    tenant_id = "1372813089198981282"
    project_id = "1372813089196881283"
    # start = datetime.fromtimestamp(1668714324) # 1666785600
    # end = start + timedelta(0, 0, 0, 1, 30, 0) # 1666799999
    start = datetime.fromtimestamp(1671054479)  # 1666785600
    end = datetime.fromtimestamp(1671055314)  # 1666799999
    proc_start = time.perf_counter()
    oom_events = asyncio.run(find_oom_evicted(cluster_id, start, end, tenant_id, project_id))
    proc_end = time.perf_counter()
    print(f"Elapsed time: {proc_end - proc_start} seconds")
    for oom_event in oom_events:
        print(
            datetime.fromtimestamp(oom_event.calibration_time)
            .replace(tzinfo=timezone.utc)
            .astimezone(timezone(timedelta(hours=8))),
            ": ",
            oom_event.payload.get("summary"),
        )