from __future__ import annotations

import asyncio

from aiohttp_retry import ExponentialRetry, RetryClient


class ClinicContext:
    def __init__(
        self,
        clinic_token,
        user_id,
        cluster_id,
        start,
        stop,
        org_id,
        org_type,
        endpoint,
        api_key,
    ):
        self.cluster_id = cluster_id
        self.clinic_token = clinic_token
        self.user_id = user_id
        self.start = start
        self.stop = stop
        self.org_id = org_id
        self.org_type = org_type
        self.endpoint = endpoint
        self.api_key = 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}",
        }
        retry_options = ExponentialRetry(attempts=5)
        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:
                print(query)
                print(
                    f"get metrics failed, status_code: {response.status}, reason: {response.reason}"
                )
            res = await response.json()
            if "status" not in res or res["status"] != "success":
                print(query)
                raise RuntimeError("PromQL was not successful.")
            if (
                "data" not in res
                or res["data"].get("resultType", None) != "matrix"
                or "result" not in res["data"]
            ):
                print(query)
                raise RuntimeError("Not valid prometheus formatted json.")
            return res["data"]["result"]

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

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