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"]
CONFIGURATION_TASK = os.environ["CONFIGURATION_TASK_NAME"]
METRICS_TASK = os.environ["METRICS_TASK_NAME"]
LOGS_TASK = os.environ["LOGS_TASK_NAME"]


class ComparationTask(BaseTask):
    cluster: str
    start_ts_1: int
    stop_ts_1: int
    start_ts_2: int
    stop_ts_2: int
    org_type: str
    org_id: str
    t1_packages: list
    t2_packages: list
    resumed: int

    def __init__(self, cluster: str, **kwargs) -> None:
        self.cluster = cluster
        self.start_ts_1 = kwargs["start_ts_1"]
        self.stop_ts_1 = kwargs["stop_ts_1"]
        self.start_ts_2 = kwargs["start_ts_2"]
        self.stop_ts_2 = kwargs["stop_ts_2"]
        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.t1_packages = (
            kwargs["package_id_1"] if kwargs.get("package_id_1") is not None else None
        )
        self.t2_packages = (
            kwargs["package_id_2"] if kwargs.get("package_id_2") is not None else None
        )
        self.resumed = 0
        self.t1_metrics_results = None
        self.t2_metrics_results = None
        self.t1_logs_results = None
        self.t2_logs_results = None
        self.t1_configuration_results = None
        self.t2_configuration_results = None
        self.t1_metrics_resp = None
        self.t2_metrics_resp = None
        self.t1_logs_resp = None
        self.t2_logs_resp = None
        self.t1_configuration_resp = None
        self.t2_configuration_resp = None
        self.context = kwargs["context"]

    @classmethod
    def resume(cls, ctx: str, lambda_context) -> ComparationTask:
        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) -> list:
        if self.org_type != "op" and self.org_type != "cloud":
            self.org_type = "op"
            # raise ValueError('org_type must be op or cloud')
        # configuration post args
        if self.org_type == "op":
            t1_configuration_post_args = {
                "args": {
                    "cluster": self.cluster,
                    "org_type": self.org_type,
                    "package_id": self.t1_packages[0],
                }
            }
            t2_configuration_post_args = {
                "args": {
                    "cluster": self.cluster,
                    "org_type": self.org_type,
                    "package_id": self.t2_packages[0],
                }
            }

        # metrics post args
        t1_metrics_post_args = {
            "args": {
                "cluster": self.cluster,
                "start_ts": self.start_ts_1,
                "stop_ts": self.stop_ts_1,
                "packages": self.t1_packages,
                "org_id": self.org_id,
                "org_type": self.org_type,
            }
        }
        t2_metrics_post_args = {
            "args": {
                "cluster": self.cluster,
                "start_ts": self.start_ts_2,
                "stop_ts": self.stop_ts_2,
                "packages": self.t2_packages,
                "org_id": self.org_id,
                "org_type": self.org_type,
            }
        }
        # logs post args
        t1_logs_post_args = {
            "args": {
                "cluster": self.cluster,
                "org_type": self.org_type,
                "start_ts": self.start_ts_1,
                "stop_ts": self.stop_ts_1,
                "org_id": self.org_id,
                "packages": self.t1_packages,
            }
        }
        t2_logs_post_args = {
            "args": {
                "cluster": self.cluster,
                "org_type": self.org_type,
                "start_ts": self.start_ts_2,
                "stop_ts": self.stop_ts_2,
                "org_id": self.org_id,
                "packages": self.t2_packages,
            }
        }
        if self.resumed == 0:
            if self.org_type == "op":
                self.t1_configuration_resp = requests.post(
                    f"https://{CE_ENDPOINT}/tasks/{CONFIGURATION_TASK}/async",
                    json.dumps(t1_configuration_post_args),
                ).json()
                self.t2_configuration_resp = requests.post(
                    f"https://{CE_ENDPOINT}/tasks/{CONFIGURATION_TASK}/async",
                    json.dumps(t2_configuration_post_args),
                ).json()
            self.t1_metrics_resp = requests.post(
                f"https://{CE_ENDPOINT}/tasks/{METRICS_TASK}/async",
                json.dumps(t1_metrics_post_args),
            ).json()
            self.t2_metrics_resp = requests.post(
                f"https://{CE_ENDPOINT}/tasks/{METRICS_TASK}/async",
                json.dumps(t2_metrics_post_args),
            ).json()
            self.t1_logs_resp = requests.post(
                f"https://{CE_ENDPOINT}/tasks/{LOGS_TASK}/async",
                json.dumps(t1_logs_post_args),
            ).json()
            self.t2_logs_resp = requests.post(
                f"https://{CE_ENDPOINT}/tasks/{LOGS_TASK}/async",
                json.dumps(t2_logs_post_args),
            ).json()
            t1_configuration_id = (
                None if self.t1_configuration_resp is None else self.t1_configuration_resp["id"]
            )
            t2_configuration_id = (
                None if self.t2_configuration_resp is None else self.t2_configuration_resp["id"]
            )
            print(
                "t1 configuration task_id: %s, t2 configuration task_id: %s, t1 metrics task_id: %s,"
                "t2 metrics task_id: %s, t1 logs task_id: %s, t2 logs task_id: %s"
                % (
                    t1_configuration_id,
                    t2_configuration_id,
                    self.t1_metrics_resp["id"],
                    self.t2_metrics_resp["id"],
                    self.t1_logs_resp["id"],
                    self.t2_logs_resp["id"],
                )
            )

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

        retry_interval = 20  # s
        max_retry = int(60 / retry_interval * 15)
        while max_retry > 0:
            if self.t1_metrics_results is None:
                t1_metrics_report = requests.get(
                    f'https://{CE_ENDPOINT}/instances/{self.t1_metrics_resp["id"]}'
                ).json()
                if (
                    t1_metrics_report["status"] == "success"
                    or t1_metrics_report["status"] == "failed"
                ):
                    self.t1_metrics_results = t1_metrics_report
                elif t1_metrics_report["status"] == "running":
                    print("wait t1 metrics results...")
                    pass

            if self.t2_metrics_results is None:
                t2_metrics_report = requests.get(
                    f'https://{CE_ENDPOINT}/instances/{self.t2_metrics_resp["id"]}'
                ).json()
                if (
                    t2_metrics_report["status"] == "success"
                    or t2_metrics_report["status"] == "failed"
                ):
                    self.t2_metrics_results = t2_metrics_report
                elif t2_metrics_report["status"] == "running":
                    print("wait t2 metrics results...")
                    pass

            if self.t1_logs_results is None:
                t1_logs_report = requests.get(
                    f'https://{CE_ENDPOINT}/instances/{self.t1_logs_resp["id"]}'
                ).json()
                if t1_logs_report["status"] == "success" or t1_logs_report["status"] == "failed":
                    self.t1_logs_results = t1_logs_report
                elif t1_logs_report["status"] == "running":
                    print("wait t1 log results...")
                    pass

            if self.t2_logs_results is None:
                t2_logs_report = requests.get(
                    f'https://{CE_ENDPOINT}/instances/{self.t2_logs_resp["id"]}'
                ).json()
                if t2_logs_report["status"] == "success" or t2_logs_report["status"] == "failed":
                    self.t2_logs_results = t2_logs_report
                elif t2_logs_report["status"] == "running":
                    print("wait t2 log results...")
                    pass

            if self.org_type == "op" and self.t1_configuration_results is None:
                t1_configuration_report = requests.get(
                    f'https://{CE_ENDPOINT}/instances/{self.t1_configuration_resp["id"]}'
                ).json()
                if (
                    t1_configuration_report["status"] == "success"
                    or t1_configuration_report["status"] == "failed"
                ):
                    self.t1_configuration_results = t1_configuration_report
                elif t1_configuration_report["status"] == "running":
                    print("wait t1 configuration results...")
                    pass

            if self.org_type == "op" and self.t2_configuration_results is None:
                t2_configuration_report = requests.get(
                    f'https://{CE_ENDPOINT}/instances/{self.t2_configuration_resp["id"]}'
                ).json()
                if (
                    t2_configuration_report["status"] == "success"
                    or t2_configuration_report["status"] == "failed"
                ):
                    self.t2_configuration_results = t2_configuration_report
                elif t2_configuration_report["status"] == "running":
                    print("wait t2 configuration 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.t1_metrics_results is not None
                and self.t2_metrics_results is not None
                and self.t1_logs_results is not None
                and self.t2_logs_results is not None
                and (
                    (
                        self.org_type == "op"
                        and self.t1_configuration_results is not None
                        and self.t2_configuration_results is not None
                    )
                    or self.org_type == "cloud"
                )
            ):
                print("comparison-report finish!")
                return [
                    self.t1_configuration_results,
                    self.t1_metrics_results,
                    self.t1_logs_results,
                    self.t2_configuration_results,
                    self.t2_metrics_results,
                    self.t2_logs_results,
                ]

    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 ComparationTask