from __future__ import annotations

import os
from pathlib import Path
from typing import Type

from context import ClinicContext
from job_avg_latency import handler as avg_latency_handler
from task import BaseTask


class LatencyTask(BaseTask):
    cluster: str
    start_ts: int
    end_ts: int
    org_id: str
    org_type: str
    endpoint: str
    api_key: str

    def __init__(
        self,
        tidb_cluster: str,
        start_ts: int,
        end_ts: int,
        org_id: str,
        org_type: str,
        **kwargs,
    ) -> None:
        self.cluster = tidb_cluster
        self.start_ts = start_ts
        self.end_ts = end_ts
        self.org_id = org_id
        self.org_type = org_type
        self.endpoint = os.environ["CLINIC_ENDPOINT"]
        self.api_key = os.environ["CLINIC_API_KEY"]

    @classmethod
    def resume(cls, ctx: str) -> LatencyTask:
        raise NotImplemented()

    async def run(self) -> None:
        async with ClinicContext(
            "token",
            "user_id",
            self.cluster,
            self.start_ts,
            self.end_ts,
            self.org_id,
            self.org_type,
            self.endpoint,
            self.api_key,
        ) as context:
            avg_latency_handler(context)

    def suspend(self) -> str:
        raise NotImplemented()


def task_class() -> Type[BaseTask]:
    return LatencyTask