from typing import Any, cast
from uuid import UUID

from pydantic import BaseModel

from ..models import (
    BrokerDetail,
    ClusterEventBatchCreating,
    ClusterEventGetIndex,
    ClusterEventIndex,
    ClusterJobsIds,
    MetricsCreating,
    MetricsLogging,
    PackageCreating,
    PackageDetail,
    RecipeCreating,
    RecipeDetail,
    ReportClusterJobResponse,
)
from .base import BaseClient


class ForCluster(BaseClient):
    class Token(BaseModel):
        token: str

    def create_package(self, body: PackageCreating) -> PackageDetail:
        res = self.request(
            "POST",
            endpoint="create-package",
            data=body,
            return_model=PackageDetail,
        )
        return cast(PackageDetail, res)

    def create_recipe(self, body: RecipeCreating) -> RecipeDetail:
        res = self.request(
            "POST",
            endpoint="create-recipe",
            data=body,
            return_model=RecipeDetail,
        )
        return cast(RecipeDetail, res)

    def update_broker(
        self,
        id: UUID,
        *,
        address: str | None = None,
        state: str | None = None,
    ) -> BrokerDetail:
        res = self.request(
            "POST",
            endpoint="update-broker",
            data={
                "id": str(id),
                "address": address,
                "state": state,
            },
            return_model=BrokerDetail,
        )
        return cast(BrokerDetail, res)

    def init_metrics(self, data: MetricsCreating) -> None:
        self.request("POST", endpoint="init-metrics", data=data)

    async def init_metrics_async(self, data: MetricsCreating) -> None:
        await self.request_async("POST", endpoint="log-metrics", data=data)

    def log_metrics(self, data: MetricsLogging) -> None:
        self.request("POST", endpoint="log-metrics", data=data)

    async def log_metrics_async(self, data: MetricsLogging) -> None:
        await self.request_async("POST", endpoint="log-metrics", data=data)

    async def report_cluster_jobs_async(
        self, data: dict[str, Any]
    ) -> ReportClusterJobResponse:
        res = await self.request_async(
            "POST",
            endpoint="report-cluster-jobs",
            data=data,
            body_model=ClusterJobsIds,
            return_model=ReportClusterJobResponse,
        )
        return cast(ReportClusterJobResponse, res)

    async def get_cluster_event_index_async(
        self, data: ClusterEventGetIndex
    ) -> ClusterEventIndex:
        res = await self.request_async(
            "POST",
            endpoint="get-cluster-event-index",
            data=data,
            return_model=ClusterEventIndex,
        )
        return cast(ClusterEventIndex, res)

    async def report_cluster_event_async(self, data: ClusterEventBatchCreating) -> None:
        await self.request_async(
            "POST",
            endpoint="report-cluster-event",
            data=data,
        )
