from __future__ import annotations

import asyncio
import json
import os

import boto3
import numpy as np
import pandas as pd
from context import ClinicContext
from utils import get_all_metrics, search_metric
from yaml import Loader, load

# from pathlib import Path


async def generate_snapshot(event) -> None:
    print(event.get("task_id"))
    print(
        f'predecessor task_id is {event.get("tags").get("source") if event.get("tags") is not None else None}'
    )
    params = event.get("kwargs")
    cluster_id = params.get("cluster")
    org_id = params.get("org_id")
    org_type = params.get("org_type").lower()
    start_ts = params.get("start_ts")
    end_ts = params.get("stop_ts")
    extra_promqls = get_extra_promqls(params)

    # Step 1: get metrics
    async with ClinicContext(cluster_id, start_ts, end_ts, org_id, org_type) as context:
        print("job start, prepare_data ...")
        template = "promql.yaml"
        all_metrics = await get_all_metrics(context, template, start_ts, end_ts, extra_promqls)

    # Step 2: do operator to metrics
    print("prepare_data done, start processing ...")
    stats = {}

    act_conn = search_metric(all_metrics, "Active_connection_count:total")
    for metrics in all_metrics:
        if metrics.status:
            metrics.init_stats()
            metrics.cal_mean()
            metrics.cal_std()
            metrics.cal_max()
            metrics.cal_min()
            # slow and lambda timeout, add if decision
            if metrics.name in [
                "QPS:total",
                "Avg_query_duration:total",
                "Active_connection_count:total",
                "Connnection_count:total",
            ]:
                metrics.cal_overall_trend()
            if "disk_usage" in metrics.name:
                metrics.cal_latest_trend(time_range=86400)
            metrics.cal_latest()
            metrics.cal_first()
            if metrics.name == "QPS:total":
                try:
                    metrics.cal_pcc(act_conn)
                except Exception as e:
                    print(e)
            if "disk_usage" in metrics.name:
                predict_time(metrics)
            if "by_type" in metrics.name:
                metrics.cal_sum()
                col = "proportion"
                tmp = []
                total_sum = np.sum(metrics.stats.loc[:, "sum"])
                for col_name, item in metrics.signals.iteritems():
                    tmp.append(metrics.stats.loc[col_name, "sum"] / (total_sum + 1e-6))
                metrics.stats.insert(0, col, tmp)
            if "uptime" in metrics.name:
                metrics.cal_falldown()
            for op in metrics.operators or []:
                eval(f"metrics.{op}")

            metrics.aggr_insight()
            metrics.cal_imbalance()  # maybe should be refactored and make aggr_insight the last one

            if metrics.rawdata == True:
                stats[metrics.name] = {
                    "unit": metrics.unit,
                    "insight": metrics.insight,
                    "rawdata": metrics.signals.to_dict(),
                }
            else:
                stats[metrics.name] = {
                    "unit": metrics.unit,
                    "insight": metrics.insight,
                }
        else:
            print(metrics.err_msg)

    result = {
        "snapshot": stats,
        "fail_metrics": {
            metrics.name: metrics.err_msg for metrics in all_metrics if not metrics.status
        },
    }
    print("processing done, start uploading ...")
    upload_res = upload_results(result, event.get("task_id"))

    return upload_res


def predict_time(metric, threshold=0.8):
    col = "remain_seconds"
    tmp = []
    threshold = threshold * 100
    for col_name, item in metric.signals.iteritems():
        if threshold <= item.tolist()[-2]:
            tmp.append(0)
        elif metric.stats.loc[col_name, "latest_trend"] <= 0:
            tmp.append(4294967295)
        else:
            tmp.append((threshold - item.tolist()[-2]) / metric.stats.loc[col_name, "latest_trend"])
    metric.stats.insert(0, col, tmp)


def upload_results(payload, job_id, store_to="s3"):
    s3 = boto3.client("s3")
    S3_BUCKET_RESPONSE = os.environ["S3_BUCKET_RESPONSE"]
    file_name = f"{job_id}.json"
    local_file = f"/tmp/{file_name}"
    folder = "metrics-snapshot"
    object_name = f"{folder}/{job_id}.json"

    with open(local_file, "w") as f:
        json.dump(payload, f, default=float)

    if store_to == "s3":
        try:
            s3.upload_file(local_file, S3_BUCKET_RESPONSE, object_name)
        except ClientError as e:
            raise ClientError(e)
        finally:
            return {"response": f"{S3_BUCKET_RESPONSE}/{object_name}"}
    else:
        raise TypeError("The type must be s3 now.")


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

        # only support yaml file now
        tmpfile = "/tmp/res.yaml"
        with open(tmpfile, "wb") as f:
            s3.download_fileobj(S3_BUCKET_REQUEST, object, f)
        with open(tmpfile, "r") as f:
            res = load(f, Loader=Loader)
        return res
    else:
        raise TypeError("The type must be s3 now.")


def get_extra_promqls(params):
    if params.get("extra_promqls") is not None:
        if params["extra_promqls"].get("from") == "s3":
            return download_payload(params["extra_promqls"]["object"])
        # default from payload
        elif params["extra_promqls"].get("from") == "payload":
            return params["extra_promqls"].get("object")
        else:
            raise TypeError("The type must be s3 or payload now.")
    else:
        return []


# TODO:
# proportion by_type anomaly
# resource higher than threshold should be implemented on business logic