from __future__ import annotations

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

import boto3
import requests
from task import BaseTask

S3_BUCKET = os.environ["S3_BUCKET"]
CE_ENDPOINT = os.environ["CE_ENDPOINT"]
LOG_DIAG_TASK = os.environ["LOG_DIAG_TASK_NAME"]
METRIC_DIAG_TASK = os.environ["METRIC_DIAG_TASK_NAME"]


class SummaryTask(BaseTask):
    cluster: str
    start_ts: int
    stop_ts: int
    org_type: str
    org_id: str
    packages: list
    resumed: int
    metrics_results: dict
    logs_results: dict
    metrics_resp: dict
    logs_resp: dict

    def __init__(self, cluster: str, start_ts: int, stop_ts: int, **kwargs) -> None:
        self.cluster = cluster
        self.start_ts = start_ts
        self.stop_ts = stop_ts
        self.org_type = kwargs["org_type"].lower() if kwargs.get("org_type") is not None else None
        self.org_id = kwargs["org_id"] if kwargs.get("org_id") is not None else None
        self.packages = kwargs["packages"] if kwargs.get("packages") is not None else None
        self.resumed = 0
        self.metrics_results = None
        self.logs_results = None
        self.metrics_resp = None
        self.logs_resp = None
        self.context = kwargs["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')
        post_args = {
            "args": {
                "cluster": self.cluster,
                "start_ts": self.start_ts,
                "stop_ts": self.stop_ts,
                "org_type": self.org_type,
                "packages": self.packages,
                "org_id": self.org_id,
            }
        }

        if self.resumed == 0:
            self.metrics_resp = requests.post(
                f"https://{CE_ENDPOINT}/tasks/{METRIC_DIAG_TASK}/async",
                json.dumps(post_args),
            ).json()
            self.logs_resp = requests.post(
                f"https://{CE_ENDPOINT}/tasks/{LOG_DIAG_TASK}/async",
                json.dumps(post_args),
            ).json()

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

        retry_interval = 20  # 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 = metrics_report["data"]
                elif metrics_report["status"] == "failed":
                    self.metrics_results = "failed"
                elif metrics_report["status"] == "running":
                    print("wait metrics results...")
                    pass

            if self.logs_results is None:
                logs_report = requests.get(
                    f'https://{CE_ENDPOINT}/instances/{self.logs_resp["id"]}'
                ).json()
                if logs_report["status"] == "success":
                    self.logs_results = logs_report["data"]
                elif logs_report["status"] == "failed":
                    self.logs_results = "failed"
                elif logs_report["status"] == "running":
                    print("wait log results...")
                    pass
            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 and self.logs_results is not None:
                if self.metrics_results != "failed":
                    for tip in self.metrics_results["advice"]:
                        if tip.get("instance") is None:
                            tip["instance"] = None
                        if tip.get("kb") is None:
                            tip["kb"] = None
                if self.logs_results != "failed":
                    for tip in self.logs_results["advice"]:
                        if tip.get("instance") is None:
                            tip["instance"] = None
                        if tip.get("kb") is None:
                            tip["kb"] = None
                if self.metrics_results != "failed" and self.logs_results != "failed":
                    summary = {
                        "metric": {
                            "advice": self.metrics_results["advice"],
                            "stats": self.metrics_results["stats"],
                            "source": "metric",
                            "timestamp": self.stop_ts,
                        },
                        "log": {
                            "advice": self.logs_results["advice"],
                            "stats": self.logs_results["stats"],
                            "source": "log",
                            "timestamp": self.stop_ts,
                        },
                    }
                if self.metrics_results == "failed" and self.logs_results != "failed":
                    summary = {
                        "metric": {},
                        "log": {
                            "advice": self.logs_results["advice"],
                            "stats": self.logs_results["stats"],
                            "source": "log",
                            "timestamp": self.stop_ts,
                        },
                    }
                if self.metrics_results != "failed" and self.logs_results == "failed":
                    summary = {
                        "metric": {
                            "advice": self.metrics_results["advice"],
                            "stats": self.metrics_results["stats"],
                            "source": "metric",
                            "timestamp": self.stop_ts,
                        },
                        "log": {},
                    }
                if self.metrics_results == "failed" and self.logs_results == "failed":
                    raise Exception("metrics and logs failed")
                print("summary finish: ", summary)
                return summary

    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