from __future__ import annotations

import os
from datetime import datetime, timezone
from typing import AsyncGenerator, Generator

from elasticsearch import AsyncElasticsearch
from elasticsearch.helpers import async_bulk
from requests import Response, Session


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.session = Session()
        self.session.trust_env = False
        self.org_id = org_id
        self.org_type = org_type
        self.endpoint = endpoint
        self.api_key = api_key

    def whoami(self):
        return self.user_id

    def get_metrics(self, query: str, start: int, end: int):
        url = "{0}/clinic/api/v1/data/metrics".format(self.endpoint)
        # url = 'http://internal-clinic-staging-private-409701995.cn-north-1.elb.amazonaws.com.cn:8085/clinic/api/v1/data/metrics'
        header = {
            "X-OrgID": self.org_id,
            "X-OrgType": self.org_type,
            "X-ClusterID": self.cluster_id,
            "Authorization": f"Bearer {self.api_key}",
        }
        response = self.session.get(
            url,
            timeout=(3, 10),
            params={
                "query": query,
                "start": str(start),
                "end": str(end),
                "step": "15s",
            },
            headers=header,
        )
        if not response.ok:
            print(
                f"get metrics failed, status_code: {response.status_code}, reason: {response.reason}"
            )
        return response

    def write_timeseries_data(self, payload: dict) -> Response:
        url = "{0}/clinic/api/v1/diagnosis/metrics_sample".format(self.endpoint)
        header = {"Authorization": f"Bearer {self.api_key}"}
        # url = 'http://internal-clinic-staging-private-409701995.cn-north-1.elb.amazonaws.com.cn:8085/clinic/api/v1/diagnosis/metrics_sample'
        response = self.session.post(url=url, timeout=(3, 10), json=payload, headers=header)
        return response