from __future__ import annotations

import asyncio
import os
from dataclasses import dataclass

from aiohttp_retry import ExponentialRetry, RetryClient


@dataclass
class ClinicContext:
    cluster_id: str
    start: int
    end: int
    org_id: str
    org_type: str
    endpoint = os.environ["CLINIC_ENDPOINT"]
    api_key = os.environ["CLINIC_API_KEY"]

    async def get_metrics(self, session: RetryClient, query: str, start: int, end: int):
        url = "{0}/clinic/api/v1/data/metrics".format(self.endpoint)
        header = {
            "X-OrgID": self.org_id,
            "X-OrgType": self.org_type,
            "X-ClusterID": self.cluster_id,
            "Authorization": f"Bearer {self.api_key}",
        }
        # query = self.replace_query(query)
        # retry_options = ExponentialRetry(attempts=2)
        # async with session.get(url, retry_options=retry_options, params={"query":query, "start": str(start), "end": str(end), "step": "15s"}, headers=header) as response:
        #     if not response.ok:
        #         result = {}
        #         result['ok'] = False
        #         result['reason'] = f"status_code: {response.status}, reason: {response.reason}, query: {query}"
        #         result['metrics'] = None
        #         return result
        #     resp = await response.json()
        #     return check_promql_resp(resp, query)
        # tricky to handle with operator deployed op cluster
        if self.org_type == "op" and isinstance(query, str) and "{label_component}" in query:
            tmp1 = query.replace("{label_cluster_id}", f'tidb_cluster="{self.cluster_id}"')
            query1 = tmp1.replace("{label_component}", "job")
            query2 = tmp1.replace("{label_component}", "component")
            retry_options = ExponentialRetry(attempts=2)
            async with session.get(
                url,
                retry_options=retry_options,
                params={
                    "query": query1,
                    "start": str(start),
                    "end": str(end),
                    "step": "15s",
                },
                headers=header,
            ) as res1:
                if not res1.ok:
                    result = {}
                    result["ok"] = False
                    result[
                        "reason"
                    ] = f"status_code: {res1.status}, reason: {res1.reason}, query: {query1}"
                    result["metrics"] = None
                resp1 = await res1.json()
            async with session.get(
                url,
                retry_options=retry_options,
                params={
                    "query": query2,
                    "start": str(start),
                    "end": str(end),
                    "step": "15s",
                },
                headers=header,
            ) as res2:
                if not res2.ok:
                    result = {}
                    result["ok"] = False
                    result[
                        "reason"
                    ] = f"status_code: {res1.status}, reason: {res1.reason}, query: {query1}"
                    result["metrics"] = None
                resp2 = await res2.json()
            if not res1.ok and not res2.ok:
                return result
            data1 = check_promql_resp(resp1, query1)
            data2 = check_promql_resp(resp2, query2)
            if data2["ok"]:
                return data2
            else:
                return data1
        else:
            query = self.replace_query(query)
            retry_options = ExponentialRetry(attempts=2)
            async with session.get(
                url,
                retry_options=retry_options,
                params={
                    "query": query,
                    "start": str(start),
                    "end": str(end),
                    "step": "15s",
                },
                headers=header,
            ) as response:
                if not response.ok:
                    result = {}
                    result["ok"] = False
                    result[
                        "reason"
                    ] = f"status_code: {response.status}, reason: {response.reason}, query: {query}"
                    result["metrics"] = None
                    return result
                resp = await response.json()
                return check_promql_resp(resp, query)

    async def __aenter__(self) -> ClinicContext:
        return self

    async def __aexit__(self, *args) -> None:
        pass

    def replace_query(self, query: str):
        if self.org_type == "cloud":
            if isinstance(query, str):
                tmp1 = query.replace("{label_cluster_id}", f'cluster_id=~".*{self.cluster_id}"')
                result = tmp1.replace("{label_component}", "component")
            if isinstance(query, dict):
                query = query.get("cloud")
                result = query.replace("{label_cluster_id}", f'cluster_id=~".*{self.cluster_id}"')
        else:
            if isinstance(query, str):
                tmp1 = query.replace("{label_cluster_id}", f'tidb_cluster="{self.cluster_id}"')
                result = tmp1.replace("{label_component}", "job")
            if isinstance(query, dict):
                query = query.get("op")
                result = query.replace("{label_cluster_id}", f'tidb_cluster="{self.cluster_id}"')
        return result


def check_promql_resp(resp, query):
    result = {}
    if "status" not in resp or resp["status"] != "success":
        result["ok"] = False
        result["reason"] = f"AIOHttp response status is not 'success'. query: {query}"
        result["metrics"] = None
        return result
    if (
        "data" not in resp
        or resp["data"].get("resultType", None) != "matrix"
        or "result" not in resp["data"]
    ):
        result["ok"] = False
        result["reason"] = f"Not valid prometheus formatted json. query: {query}"
        result["metrics"] = None
        return result
    if len(resp["data"]["result"]) == 0:
        result["ok"] = False
        result["reason"] = f"PromQL returns null. query: {query}"
        result["metrics"] = None
        return result
    if len(resp["data"]["result"]) == 1 and resp["data"]["result"][0]["metric"] == {}:
        result["ok"] = False
        result["reason"] = f"PromQL is not right. query: {query}"
        result["metrics"] = None
        return result
    result["ok"] = True
    result["reason"] = f"OK. query: {query}"
    result["metrics"] = resp["data"]["result"]
    return result