import json
import os
import time
from datetime import date, datetime

import numpy as np
import pandas as pd
import requests

CLINIC_ENDPOINT = os.environ["CLINIC_ENDPOINT"]
CLINIC_API_KEY = os.environ["CLINIC_API_KEY"]


def lambda_handler(event, context):
    print(event)
    args = event.get("kwargs") if event.get("kwargs") is not None else {}
    _date = (
        datetime.strptime(args.get("date"), "%Y-%m-%d").date()
        if args.get("date") is not None
        else date.today()
    )
    thresholds = args.get("threshold") if args.get("threshold") is not None else [1, 100, 10000]
    clusters = args.get("clusters")
    res = []

    df = sync_metrics(_date, clusters)
    tenant_dict = get_tenant_tricky(clusters)
    if df.shape[0] > 0:
        df["time"] = pd.to_datetime(df["time"], unit="s")
        df["value"] = df["value"].astype(float)
        for threshold in thresholds:
            res.extend(gt_threshold(df, threshold, tenant_dict))
    result = write_data(res)
    return result


def gt_threshold(data, threshold, tenant_dict):
    df = data[(data.value > threshold)]
    df = df.pivot(index="time", columns="cluster_id", values="value")
    rows = []
    for col_name, qps in df.iteritems():
        time_ranges = get_timerange_list(qps)
        for timerange in time_ranges:
            rows.append(
                {
                    "cluster_id": modify_cluster_id(col_name),
                    "org_id": tenant_dict[col_name],
                    "start_unix": int(df.index[timerange[0]].timestamp()),
                    "end_unix": int(df.index[timerange[1] - 1].timestamp() + 60),
                    "max_qps": max(qps[timerange[0] : timerange[1]]),
                    "threshold": threshold,
                }
            )
    return rows


def cal_continuous_flag_slice(flag_slice, continuous_num=1):
    rising_edge = []
    falling_edge = []
    res = []
    for i in range(len(flag_slice) - 1):
        if flag_slice[i] == 0 and flag_slice[i + 1] == 1:
            rising_edge.append(i)
        if flag_slice[i] == 1 and flag_slice[i + 1] == 0:
            falling_edge.append(i)
    i = len(rising_edge)
    j = len(falling_edge)
    if i == 0 and j == 0:
        if flag_slice[0] == 1:
            res.append((0, len(flag_slice)))
    elif i == 0 and j == 1:
        res.append((0, falling_edge[0] + 1))
    elif i == 1 and j == 0:
        res.append((rising_edge[0] + 1, len(flag_slice)))
    elif rising_edge[0] < falling_edge[0]:  # rising fisrt
        for s in range(j):
            if falling_edge[s] - rising_edge[s] >= continuous_num:
                res.append((rising_edge[s] + 1, falling_edge[s] + 1))
        if j < i:
            if len(flag_slice) - rising_edge[-1] >= continuous_num:
                res.append((rising_edge[-1] + 1, len(flag_slice)))
    elif rising_edge[0] > falling_edge[0]:  # falling first
        if falling_edge[0] + 1 >= continuous_num:
            res.append((0, falling_edge[0] + 1))
        for s in range(j - 1):
            if falling_edge[s + 1] - rising_edge[s] >= continuous_num:
                res.append((rising_edge[s] + 1, falling_edge[s + 1] + 1))
        if i == j:
            if len(flag_slice) - rising_edge[-1] >= continuous_num:
                res.append((rising_edge[-1] + 1, len(flag_slice)))
    return res


def get_timerange_list(alist):
    sign_list = [0 if np.isnan(item) else 1 for item in alist]
    return cal_continuous_flag_slice(sign_list)


def modify_cluster_id(cluster_id):
    if "tidb" in cluster_id:
        return cluster_id.replace("tidb", "")
    else:
        return cluster_id


def sync_metrics(_date, clusters):
    start = int(time.mktime(_date.timetuple()))
    end = start + 24 * 60 * 60 - 1

    url = f"{CLINIC_ENDPOINT}/clinic/api/v1/data/metrics"
    r = {}
    df_total = pd.DataFrame()
    for cluster in clusters:
        cluster_id = cluster.get("cluster_id")
        header = {
            "X-OrgID": cluster.get("org_id"),
            "X-OrgType": "cloud",
            "X-ClusterID": cluster_id,
            "Authorization": f"Bearer {CLINIC_API_KEY}",
        }
        query = 'sum(rate(tidb_executor_statement_total{cluster_id=~".*{clusterid}"}[1m]))'.replace(
            "{clusterid}", cluster_id
        )
        response = requests.get(
            url,
            params={"query": query, "start": start, "end": end, "step": "1m"},
            headers=header,
        )
        if response.ok:
            content = json.loads(response.text)
            if content.get("data") is not None and content["data"].get("result") is not None:
                if len(content["data"]["result"]) > 0:
                    df = pd.DataFrame.from_records(
                        [
                            {
                                "cluster_id": cluster_id,
                                "time": x[0],
                                "value": x[1],
                            }
                            for x in content["data"]["result"][0]["values"]
                        ]
                    )
                    df_total = pd.concat([df_total, df], ignore_index=True)
        else:
            print(f'promql "{query}"request failed with {response.status_code}-{response.reason}')
    return df_total


def get_tenant_tricky(clusters):
    tenant_dict = {}
    for cluster in clusters:
        tenant_dict[cluster.get("cluster_id")] = cluster.get("org_id")
    return tenant_dict


def write_data(res):
    url = f"{CLINIC_ENDPOINT}/clinic/api/v1/data/active_qps/batch_insert"
    header = {
        "Authorization": f"Bearer {CLINIC_API_KEY}",
        "Content-Type": "application/json",
    }
    body = {"items": res}
    print(body)
    if len(res) > 0:
        response = requests.post(url, data=json.dumps(body, default=float), headers=header)
        content = json.loads(response.text)
        print(content)
        if content.get("status") == "ok":
            return {"statusCode": 200, "length": len(res)}
        elif content.get("error"):
            return {
                "statusCode": 500,
                "message": content.get("message"),
                "length": len(res),
            }
        else:
            return {
                "statusCode": 500,
                "message": "undefined",
                "length": len(res),
            }
    else:
        return {"statusCode": 200, "items": []}