from __future__ import annotations

import json
import os
import pickle
import time
import uuid
from pathlib import Path
from typing import Type

import boto3
import numpy as np
import requests
from mathbox.calculus.differential import diff
from task import BaseTask
from yaml import Loader, load

S3_BUCKET = os.environ["S3_BUCKET"]
CE_ENDPOINT = os.environ["CE_ENDPOINT"]
METRICS_SNAPSHOT_TASK_NAME = os.environ["METRICS_SNAPSHOT_TASK_NAME"]


class SummaryTask(BaseTask):
    task_id: str
    cluster: str
    start_ts: int
    stop_ts: int
    org_type: str
    org_id: str
    resumed: int
    metrics_results: dict
    metrics_resp: dict

    def __init__(self, event) -> None:
        print(
            f'predecessor task_id is {event.get("tags").get("source") if event.get("tags") is not None else None}'
        )
        self.task_id = event.get("task_id") or uuid.uuid4()
        params = event.get("kwargs")
        self.cluster = params.get("cluster")
        print(f"cluster id is {self.cluster}")
        self.start_ts = params.get("start_ts")
        self.stop_ts = params.get("stop_ts")
        self.org_type = params.get("org_type").lower()
        self.org_id = params.get("org_id")
        self.resumed = 0
        self.extra_promqls = {
            "from": "payload",
            "object": [
                {
                    "name": "TiKV_safepoint:total",
                    "unit": "s",
                    "expr": 'label_replace(max(tikv_gcworker_autogc_safe_point{{label_cluster_id}})/(2^18),"name","tikv_safepoint","","")',
                    "rawdata": True,
                    "operators": [],
                }
            ],
        }
        self.metrics_results = None
        self.metrics_resp = None
        self.context = params.get("context")

    @classmethod
    def resume(cls, ctx: str, lambda_context) -> SummaryTask:
        s3_object = ctx
        s3 = boto3.client("s3")
        tmpfile = "/tmp/resume.pkl"
        with open(tmpfile, "wb") as f:
            s3.download_fileobj(S3_BUCKET, s3_object, f)
        task = pickle.load(open(tmpfile, "rb"))
        task.resumed += 1
        task.context = lambda_context
        print(f"resumed {task.resumed} times.")
        return task

    async def run(self) -> dict:
        if self.org_type != "op" and self.org_type != "cloud":
            self.org_type = "op"
            # raise ValueError('org_type must be op or cloud')

        p = Path(__file__).with_name("rules.yaml")
        with p.open("r") as f:
            rules = load(f, Loader=Loader)
        for rule in rules:
            if rule["switch"]:
                for metric in rule["metrics"]:
                    if metric.get("expr") is not None:
                        self.extra_promqls["object"].append(
                            {
                                "name": metric["name"],
                                "unit": metric["unit"],
                                "expr": metric["expr"],
                                "rawdata": metric.get("rawdata"),
                                "operators": metric.get("operators"),
                            }
                        )

        post_args = {
            "tags": {"source": self.task_id},
            "args": {
                "cluster": self.cluster,
                "start_ts": self.start_ts,
                "stop_ts": self.stop_ts,
                "org_type": self.org_type,
                "org_id": self.org_id,
                "extra_promqls": self.extra_promqls,
            },
        }

        if self.resumed == 0:
            # Step 1: get snapshot
            print("Getting the snapshot...")
            self.metrics_resp = requests.post(
                f"https://{CE_ENDPOINT}/tasks/{METRICS_SNAPSHOT_TASK_NAME}/async",
                json.dumps(post_args),
            ).json()
            print("snapshot task id is ", self.metrics_resp)

        if self.resumed >= 20:
            raise Exception("diag-summary-op-test runs over 150min, stopped.")

        retry_interval = 5  # s
        max_retry = int(60 / retry_interval * 15)
        while max_retry > 0:
            if self.metrics_results is None:
                metrics_report = requests.get(
                    f'https://{CE_ENDPOINT}/instances/{self.metrics_resp["id"]}'
                ).json()
                if metrics_report["status"] == "success":
                    self.metrics_results = download_result(
                        f'{METRICS_SNAPSHOT_TASK_NAME}/{self.metrics_resp["id"]}.json'
                    )
                    # self.metrics_results = metrics_report['data']
                elif metrics_report["status"] == "failed":
                    self.metrics_results = {
                        "status": "failed",
                        "errMsg": metrics_report["data"],
                    }
                elif metrics_report["status"] == "running":
                    print("wait metrics results...")

            if self.context.get_remaining_time_in_millis() / 1000 < (retry_interval * 2 + 1):
                ctx = self.suspend()
                print("continuation context: ", ctx)
                return {
                    "code": 206,
                    "context": ctx,
                }
            time.sleep(retry_interval)
            max_retry -= 1
            if self.metrics_results is not None:
                if self.metrics_results.get("snapshot") is not None:
                    stats = self.metrics_results.get("snapshot")
                    fail_metrics = self.metrics_results.get("fail_metrics")

                    if stats == {}:
                        raise Exception("no metrics data")

                    # print(self.metrics_results)

                    # Step 2: other data process on snapshot
                    for k, v in stats.items():
                        if "by_instance" in k:
                            insight_mean = {
                                kk: val if val is not None else 0
                                for kk, val in v["insight"]["mean"].items()
                            }
                            insight_max = {
                                kk: val if val is not None else 0
                                for kk, val in v["insight"]["max"].items()
                            }
                            insight_min = {
                                kk: val if val is not None else 0
                                for kk, val in v["insight"]["min"].items()
                            }
                            stats[k]["insight"]["max_kv_avg"] = (
                                max(insight_mean, key=insight_mean.get),
                                max(insight_mean.values()),
                            )
                            stats[k]["insight"]["min_kv_avg"] = (
                                min(insight_mean, key=insight_mean.get),
                                min(insight_mean.values()),
                            )
                            stats[k]["insight"]["mean_mean"] = np.mean(list(insight_mean.values()))
                            stats[k]["insight"]["max_max"] = (
                                max(insight_max, key=insight_max.get),
                                max(insight_max.values()),
                            )  # np.max(list(v['insight']['max'].values()))
                            stats[k]["insight"]["min_min"] = (
                                min(insight_min, key=insight_min.get),
                                min(insight_min.values()),
                            )  # np.min(list(v['insight']['min'].values()))

                    if stats.get("Database_time:total") is not None:
                        if stats["Database_time:total"]["insight"].get("mean") is not None:
                            get_token_proportion = (
                                stats["Get_token_time:total"]["insight"]["mean"]
                                / (stats["Database_time:total"]["insight"]["mean"] + 1e-6)
                                * 100
                                if stats.get("Get_token_time:total") is not None
                                and stats.get("Database_time:total") is not None
                                else None
                            )  # avoid zero division
                            parse_proportion = (
                                stats["Parse_time:total"]["insight"]["mean"]
                                / (stats["Database_time:total"]["insight"]["mean"] + 1e-6)
                                * 100
                                if stats.get("Parse_time:total") is not None
                                and stats.get("Database_time:total") is not None
                                else None
                            )  # avoid zero division
                            compile_proportion = (
                                stats["Compile_time:total"]["insight"]["mean"]
                                / (stats["Database_time:total"]["insight"]["mean"] + 1e-6)
                                * 100
                                if stats.get("Compile_time:total") is not None
                                and stats.get("Database_time:total") is not None
                                else None
                            )  # avoid zero division
                            execute_proportion = (
                                stats["Execute_time:total"]["insight"]["mean"]
                                / (stats["Database_time:total"]["insight"]["mean"] + 1e-6)
                                * 100
                                if stats.get("Execute_time:total") is not None
                                and stats.get("Database_time:total") is not None
                                else None
                            )  # avoid zero division
                            stats["Database_time:by_phase"] = {
                                "unit": "%",
                                "insight": {
                                    "proportion": {
                                        "Get_token_time:total": get_token_proportion,
                                        "Parse_time:total": parse_proportion,
                                        "Compile_time:total": compile_proportion,
                                        "Execute_time:total": execute_proportion,
                                    }
                                },
                            }

                    if stats.get("TiKV_safepoint:total") is not None:
                        if stats["TiKV_safepoint:total"].get("rawdata") is not None:
                            safepoint_timestamp = [
                                int(a)
                                for a in list(
                                    stats["TiKV_safepoint:total"]["rawdata"][
                                        "tikv_safepoint"
                                    ].keys()
                                )
                            ]
                            safepoint_curve = [
                                float(a)
                                for a in list(
                                    stats["TiKV_safepoint:total"]["rawdata"][
                                        "tikv_safepoint"
                                    ].values()
                                )
                            ]
                            safepoint_changetime = [safepoint_timestamp[0]]
                            for i in range(len(safepoint_curve) - 1):
                                if safepoint_curve[i + 1] - safepoint_curve[i] > 0:
                                    safepoint_changetime.append(safepoint_timestamp[i + 1])
                            if len(safepoint_changetime) > 1:
                                diff_timestamp = diff(safepoint_changetime)
                            else:
                                gc_lifetime = 600  # default config
                                diff_timestamp = [
                                    0,
                                    gc_lifetime,
                                    (self.stop_ts - self.start_ts),
                                ]
                            stats["TiKV_safepoint:total"]["insight"]["interval"] = diff_timestamp
                    # Step 3: analyze rule yaml
                    print("Checking diagnosis rules...")

                    advice = []
                    pass_rule_cnt = 0
                    fail_rule_cnt = 0
                    error_rule_cnt = 0

                    # p = Path(__file__).with_name('rules.yaml')
                    # with p.open('r') as f:
                    # rules = load(f, Loader=Loader)

                    for rule in rules:
                        if rule["switch"]:
                            rule_lack_varible = False
                            for item in rule["metrics"]:
                                metric = stats.get(item["name"])
                                if metric is None:
                                    print(f'rule {rule["id"]}: {item["name"]} has nodata')
                                    rule_lack_varible = True
                                    break
                                else:
                                    for k, v in item.get("variables").items():
                                        fields = v.split(".")
                                        locals()[k] = get_stats(metric.get("insight"), fields)
                            if rule_lack_varible:
                                error_rule_cnt += 1
                                continue
                            try:
                                if eval(rule["condition"]):
                                    if rule["end2end"]:
                                        label = "End-to-End"
                                    else:
                                        label = "Non-End-to-End"
                                    variables = {}
                                    piece = {
                                        "id": rule["id"],
                                        "name": rule["name"],
                                        "qualified": rule["qualified"],
                                        "level": rule["level"],
                                        "category": rule["category"],
                                        "component": rule["component"],
                                        "metrics": [x["name"] for x in rule["metrics"]],
                                        "type": label,
                                        "content": {
                                            "key": rule["advice"]["key"],
                                            "variables": None,
                                        },
                                    }
                                    if "Database_time:by_phase" in piece["metrics"]:
                                        if (
                                            "Get_token_time:total"
                                            in list(rule["metrics"][0]["variables"].values())[0]
                                        ):
                                            piece["metrics"] = [
                                                "Database_time:total",
                                                "Get_token_time:total",
                                            ]
                                        elif (
                                            "Parse_time:total"
                                            in list(rule["metrics"][0]["variables"].values())[0]
                                        ):
                                            piece["metrics"] = [
                                                "Database_time:total",
                                                "Parse_time:total",
                                            ]
                                        elif (
                                            "Compile_time:total"
                                            in list(rule["metrics"][0]["variables"].values())[0]
                                        ):
                                            piece["metrics"] = [
                                                "Database_time:total",
                                                "Compile_time:total",
                                            ]
                                    if rule["advice"]["variables"] is not None:
                                        for variable in rule["advice"]["variables"]:
                                            ((k, v),) = variable.items()
                                            try:
                                                v = eval(v)
                                            except Exception as e:
                                                print(e)
                                            if isinstance(v, float):
                                                variables[k] = f"{v:.2f}"
                                            else:
                                                variables[k] = f"{v}"
                                        piece["content"]["variables"] = variables
                                    advice.append(piece)
                                    fail_rule_cnt += 1
                                else:
                                    pass_rule_cnt += 1
                            except Exception as e:
                                print(f'rule {rule["id"]} exception: {e}')
                                error_rule_cnt += 1
                                continue

                    rule_stats = {
                        "passedRules": pass_rule_cnt,
                        "failedRules": fail_rule_cnt,
                        "errorRules": error_rule_cnt,
                        #'scannedMetrics': len(all_metrics)
                    }

                    metrics_stats = []
                    p = Path(__file__).with_name("report.yaml")
                    with p.open("r") as f:
                        metrics = load(f, Loader=Loader)
                    for metric in metrics:
                        category = []
                        for item in metric.get("category"):
                            category.append(
                                {
                                    "name": list(item.keys())[0],
                                    "main": list(item.values())[0],
                                }
                            )
                        tmp = {
                            "name": metric.get("name"),
                            "category": category,
                            "values": {},
                            "errMsg": None,
                        }
                        tmp["errMsg"] = fail_metrics.get(metric.get("name"))
                        if tmp["errMsg"] is not None:
                            metrics_stats.append(tmp)
                            continue
                        unit = stats[metric.get("name")]["unit"]
                        for value in metric.get("variables"):
                            orig_value = stats[metric.get("name")]["insight"][value]
                            if isinstance(orig_value, dict):
                                if value == "proportion" and unit != "%":
                                    orig_value = topk(orig_value, 3)
                                for k, v in orig_value.items():
                                    orig_value[k] = adapt_unit(v, unit)
                                    if value == "proportion" and unit != "%":
                                        orig_value[k] = adapt_unit(v * 100, "%")
                                tmp["values"][value] = orig_value
                            elif isinstance(orig_value, tuple):
                                tmp["values"][value] = (
                                    orig_value[0],
                                    adapt_unit(orig_value[1], unit),
                                )
                            else:
                                tmp["values"][value] = adapt_unit(orig_value, unit)
                            if tmp["values"].get("mean_mean"):
                                tmp["values"]["mean"] = tmp["values"].pop("mean_mean")
                            if tmp["values"].get("max_max"):
                                tmp["values"]["max"] = tmp["values"].pop("max_max")
                            if tmp["values"].get("min_min"):
                                tmp["values"]["min"] = tmp["values"].pop("min_min")
                        metrics_stats.append(tmp)

                    return {
                        "abnormal": len(advice) > 0,
                        "advice": advice,
                        "ruleStats": rule_stats,
                        "metricStats": metrics_stats,
                    }
                else:
                    return {"err_msg": f'snapshot task {self.metrics_resp["id"]} failed'}

    def suspend(self) -> str:
        s3_object = f"summary-{uuid.uuid4()}.pkl"
        s3 = boto3.client("s3")
        with open(f"/tmp/{s3_object}", "wb") as handle:
            pickle.dump(self, handle)
        try:
            s3.upload_file(f"/tmp/{s3_object}", S3_BUCKET, s3_object)
        except Exception as e:
            print(e)
            return False
        return s3_object


def task_class() -> Type[BaseTask]:
    return SummaryTask


def get_stats(metric, fields: list):
    if len(fields) == 0:
        return metric
    tmp = metric.get(fields[0])
    return get_stats(tmp, fields[1:])


def adapt_unit(value, unit):
    if value is None:
        return None
    if unit == "1":
        return "%.2f" % value
    elif unit == "s":
        units = ["s", "ms", "us"]
        size = 1000.0
        for i in range(len(units)):
            if value > 1 or units[i] == "us":
                return "%.2f%s" % (value, units[i])
            value = value * size
    elif unit == "Bytes":
        units = ["B", "KB", "MB", "GB", "TB", "PB"]
        size = 1024.0
        for i in range(len(units)):
            if value / size < 1:
                return "%.2f%s" % (value, units[i])
            value = value / size
    elif unit == "%":
        return "%.2f%%" % (value)
    elif unit == "Bps":
        units = ["B/s", "KB/s", "MB/s", "GB/s", "TB/s", "PB/s"]
        size = 1024.0
        for i in range(len(units)):
            if value / size < 1:
                return "%.2f%s" % (value, units[i])
            value = value / size


def download_result(object, get_from="s3"):
    if get_from == "s3":
        s3 = boto3.client("s3")
        S3_BUCKET_RESPONSE = os.environ["S3_BUCKET_RESPONSE"]

        tmpfile = "/tmp/res.json"
        with open(tmpfile, "wb") as f:
            s3.download_fileobj(S3_BUCKET_RESPONSE, object, f)
        with open(tmpfile, "r") as f:
            res = json.load(f)
        return res
    else:
        raise TypeError("The type must be s3 now.")


def topk(d, k):
    return dict(sorted(d.items(), key=lambda item: item[1], reverse=True)[:k])