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 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"]
EXTRA_PROMQLS = os.environ["REQUEST_EXTRA_PROMQL"]


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(event)
        self.task_id = event.get("task_id")  # or uuid.uuid4()
        params = event.get("kwargs")
        self.cluster = params.get("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 = []
        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('extra_promql.yaml')
        # with p.open('r') as f:
        #    self.extra_promqls = load(f, Loader=Loader)
        # self.extra_promqls = upload_request('extra_promql.yaml', self.task_id)
        self.extra_promqls = {"from": "s3", "object": EXTRA_PROMQLS}

        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") or {}

                    # print(self.metrics_results)

                    # Step 2: other data process on snapshot
                    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
                            if execute_proportion > 100:
                                execute_proportion = 100
                            stats["Database_time:by_phase"] = {
                                "unit": "s",
                                "insight": {
                                    "proportion": {
                                        "unit": "%",
                                        "value": {
                                            "Get_token_time:total": get_token_proportion,
                                            "Parse_time:total": parse_proportion,
                                            "Compile_time:total": compile_proportion,
                                            "Execute_time:total": execute_proportion,
                                        },
                                    }
                                },
                            }

                    # Step 3: orgnize the report

                    metrics_stats = []
                    # p = Path(__file__).with_name('output.yaml')
                    # with p.open('r') as f:
                    #     metrics = load(f, Loader=Loader)
                    # for metric in metrics:
                    #     tmp = {
                    #         'name': metric.get('name'),
                    #         '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):
                    #             for k,v in orig_value.items():
                    #                 orig_value[k] = adapt_unit(v, unit)
                    #                 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)
                    #     metrics_stats.append(tmp)
                    for name, insight in stats.items():
                        tmp = {"name": name, "unit": "", "values": {}, "errMsg": None}
                        tmp["errMsg"] = fail_metrics.get(name)
                        if tmp["errMsg"] is not None:
                            metrics_stats.append(tmp)
                            continue
                        unit = stats[name]["unit"]
                        if name == "Database_time:by_phase":
                            orig_value = stats[name]["insight"]["proportion"]
                            #    for k,v in orig_value.items():
                            #        orig_value[k] = adapt_unit(v, unit)
                            #        tmp['values']['proportion'] = orig_value
                            tmp["unit"] = unit
                            tmp["values"]["proportion"] = orig_value
                            metrics_stats.append(tmp)
                            continue
                        if name == "DDL_owner:by_instance":
                            start_kv = []
                            end_kv = []
                            for k, v in stats[name]["rawdata"].items():
                                v = list(v.items())
                                start_kv.append((k, v[0]))
                                end_kv.append((k, v[-2]))
                            try:
                                ddl_owner_start = start_kv[
                                    np.nanargmax([x[1][1] for x in start_kv])
                                ][0]
                            except ValueError:
                                ddl_owner_start = None
                            try:
                                ddl_owner_end = end_kv[np.nanargmax([x[1][1] for x in end_kv])][0]
                            except ValueError:
                                ddl_owner_end = None
                            tmp["values"]["ddl_owner_start"] = {"value": ddl_owner_start}
                            tmp["values"]["ddl_owner_end"] = {"value": ddl_owner_end}
                        if name == "PD_leader:by_instance":
                            start_kv = []
                            end_kv = []
                            for k, v in stats[name]["rawdata"].items():
                                v = list(v.items())
                                start_kv.append((k, v[0]))
                                end_kv.append((k, v[-2]))
                            try:
                                pd_leader_start = start_kv[
                                    np.nanargmax([x[1][1] for x in start_kv])
                                ][0]
                            except ValueError:
                                pd_leader_start = None
                            try:
                                pd_leader_end = end_kv[np.nanargmax([x[1][1] for x in end_kv])][0]
                            except ValueError:
                                pd_leader_end = None
                            tmp["values"]["pd_leader_start"] = {"value": pd_leader_start}
                            tmp["values"]["pd_leader_end"] = {"value": pd_leader_end}
                        if name == "Component:total" or name == "Connnection_count:total":
                            # first and last
                            tmp["unit"] = "1"
                            tmp["values"]["first"] = {"value": stats[name]["insight"]["first"]}
                            tmp["values"]["latest"] = {"value": stats[name]["insight"]["latest"]}
                        if (
                            name == "TiFlash_count:total"
                            or name == "TiKV_count:total"
                            or name == "TiDB_count:total"
                            or name == "PD_count:total"
                        ):
                            tmp["unit"] = "1"
                            tmp["values"]["max"] = {"value": stats[name]["insight"]["max"]}
                        for variable in ["mean", "std"]:
                            orig_value = stats[name]["insight"][variable]
                            if isinstance(orig_value, dict):
                                # for k,v in orig_value.items():
                                #    orig_value[k] = adapt_unit(v, unit)
                                #    tmp['values'][variable] = orig_value
                                tmp["unit"] = unit
                                tmp["values"][variable] = {"value": orig_value}
                            elif isinstance(orig_value, tuple):
                                # tmp['values'][variable] = (orig_value[0], adapt_unit(orig_value[1], unit))
                                tmp["unit"] = unit
                                tmp["values"][variable] = {
                                    "pos": orig_value[0],
                                    "value": orig_value[1],
                                }
                            else:
                                # tmp['values'][variable] = adapt_unit(orig_value, unit)
                                tmp["unit"] = unit
                                tmp["values"][variable] = {"value": orig_value}
                        if name == "Database_time:by_type" or name == "SQL_execute_time:by_type":
                            proportion = stats[name]["insight"]["proportion"]
                            proportion.update((k, v * 100) for k, v in proportion.items())
                            tmp["values"]["proportion"] = {"unit": "%", "value": proportion}
                        metrics_stats.append(tmp)
                    new_dict = replace_nan_with_none(metrics_stats)
                    return {"metricStats": new_dict, "failedMetrics": list(fail_metrics.keys())}
                else:
                    return {
                        "metricStats": None,
                        "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 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 upload_request(file_name, job_id, store_to="s3"):
    s3 = boto3.client("s3")
    S3_BUCKET_REQUEST = os.environ["S3_BUCKET_REQUEST"]
    local_file = str(Path(__file__).with_name(file_name))
    folder = "metrics-snapshot"
    object_name = f"{folder}/{job_id}.json"

    if store_to == "s3":
        try:
            s3.upload_file(local_file, S3_BUCKET_REQUEST, object_name)
        except Exception as e:
            print(e)
        finally:
            return {"from": "s3", "object": f"{object_name}"}
    else:
        raise TypeError("The type must be s3 now.")


def replace_nan_with_none(d):
    if isinstance(d, dict):
        for key in d:
            d[key] = replace_nan_with_none(d[key])
    elif isinstance(d, list):
        for i, item in enumerate(d):
            d[i] = replace_nan_with_none(item)
    elif isinstance(d, float) and np.isnan(d):
        return None
    return d