import asyncio
import time
from dataclasses import dataclass
from pathlib import Path
from typing import List

import aiohttp
import numpy as np
import pandas as pd
from aiohttp_retry import ExponentialRetry, RetryClient
from context import ClinicContext

# from mathbox.app.signal.filter import moving_median, f_lowpass_filter
# from mathbox.app.signal.outlier import noise_outlier, simple_outlier
# from mathbox.na.interpolation import series_linear_interp
# from mathbox.statistics.estimator import mean, std, median
from yaml import Loader, load


def timer(func):
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        res = func(*args, **kwargs)
        end = time.perf_counter()
        print(f"Running time：{(end-start):.3f} s")
        return res

    return wrapper


def async_timer(func):
    async def process(func, *args, **params):
        if asyncio.iscoroutinefunction(func):
            return await func(*args, **params)
        else:
            print("this is not a coroutine")
            return func(*args, **params)

    async def helper(*args, **params):
        start = time.perf_counter()
        result = await process(func, *args, **params)
        print("Running time：", time.perf_counter() - start, "s")
        return result

    return helper


@dataclass
class Metric:
    # One PromQL query with name, timestamp, query response and signals
    id: int
    name: str
    unit: str
    query: str
    component: str
    type: list
    sub_metrics: list
    only_sub: bool
    description: str
    panel: str
    stats = {}
    signals = pd.DataFrame()  # list of Signal


# @async_timer
async def preprocess(session, metric: Metric, context, desired_timestamp):
    """
    Query PromQL and write the result,which is not imputed as pandas.DataFrame into metric.signals.
    """
    try:
        res = await context.get_metrics(
            session, metric.query, context.start, context.stop
        )  # list of dicts
    except Exception as e:
        print(e)
        res = []
    step = 15
    timestamp = np.arange(context.start, context.stop + step, step)
    cols = len(res) if len(res) > 0 else 1

    data = np.zeros((cols, len(timestamp)), dtype=np.float32)
    # data[:] = np.nan # prometheus is rediculous, request data is NaN, but treats as zero.
    col_names = []

    if len(res) > 0:
        for i, sig in enumerate(res):
            # sig: {"metric":{"name": "tidb-01"}, "values":[[t,v],[t,v],...,[t,v]]}
            if sig["metric"].get("name") is not None:
                sig_name = sig["metric"]["name"]
            else:
                sig_name = f'{sig["metric"].get("name1")}-{sig["metric"].get("name2")}-{sig["metric"].get("name3")}'
            t, v = zip(*sig["values"])
            index = np.rint((np.array(t) - context.start) / step).astype(np.int)
            data[i, index] = [float(val) for val in v]
            col_names.append(sig_name)
        metric.signals = pd.DataFrame(
            data.T,
            columns=col_names,
            index=pd.to_datetime(timestamp, unit="s"),
        )
        metric.signals.fillna(0, inplace=True)

        # step 1: check if too many NaNs, too many raise an errror
        # metric.signals.interpolate(method='index', inplace=True, axis=0, limit_direction='both') # step 2: imputation
    else:
        metric.signals = pd.DataFrame(
            data.T,
            columns=[metric.name],
            index=pd.to_datetime(timestamp, unit="s"),
        )
        metric.signals.replace(0, np.nan, inplace=True)

    return metric


async def get_all_metrics(context: ClinicContext, fname: str, start: int, end: int):
    p = Path(__file__).with_name(fname)
    with p.open("r") as f:
        metrics_req = load(f, Loader=Loader)

        desired_timestamp = list(range(start, end + 15, 15))

        promqls = []
        for item in metrics_req:
            query = item["expr"].replace("{tidb_cluster_id}", context.cluster_id)
            if item.get("sub_metrics") == None:
                sub_metrics = []
            else:
                sub_metrics = item["sub_metrics"]
            if item.get("only_sub") == None:
                only_sub = False
            else:
                only_sub = item["only_sub"]
            promqls.append(
                Metric(
                    item["id"],
                    item["name"],
                    item["unit"],
                    query,
                    item["component"],
                    item["type"],
                    sub_metrics,
                    only_sub,
                    item["description"],
                    item["panel"],
                )
            )

        conn = aiohttp.TCPConnector(limit=2)
        # retry_options = ExponentialRetry(attempts=5)
        # async with RetryClient(raise_for_status=False, retry_options=retry_options, connector=conn) as session:
        async with RetryClient(raise_for_status=False, connector=conn) as session:
            metrics = await asyncio.gather(
                *[preprocess(session, item, context, desired_timestamp) for item in promqls],
                return_exceptions=True,
            )

        return metrics


# not adjust to new metric class with dataframe
def cal_imbalance(pieces, ratio_threshold=0.5):
    ratio = 1 - min(pieces) / (max(pieces) + 1e-10)  # not break change but important
    if ratio >= ratio_threshold:
        return True
    else:
        return False


# not adjust to new metric class with dataframe
def quantile(data, q):
    data.sort()
    return data[int(len(data) * q)]