from dataclasses import dataclass
from typing import Dict

from ydb.public.api.grpc.draft.ydb_persqueue_v1_pb2_grpc import ClusterDiscoveryServiceStub
from ydb.public.api.protos.ydb_persqueue_cluster_discovery_pb2 import(
    DiscoverClustersRequest, DiscoverClustersResult, ReadSessionParams, WriteSessionParams
)

from saas.library.python.logbroker.client import LogbrokerClient, TServiceClass


@dataclass
class ClusterInfo:
    name: str
    endpoint: str
    available: bool


@dataclass
class ClustersInfo:
    dc_to_cluster_info: Dict[str, ClusterInfo]


class ClusterInfoClient(LogbrokerClient):
    _PROTO_MAPPING = {
        'DiscoverClusters': DiscoverClustersRequest
    }

    @property
    def _service_stub_cls(self) -> TServiceClass:
        return ClusterDiscoveryServiceStub

    @staticmethod
    def _write_dc_availability(data: dict, sessions) -> None:
        for session in sessions:
            for cluster in session.clusters:
                curr_cluster_info = ClusterInfo(cluster.name, cluster.endpoint, cluster.available)

                last_cluster_info = data.get(cluster.name)
                if last_cluster_info:
                    curr_cluster_info.available = last_cluster_info.available and curr_cluster_info.available

                data[cluster.name] = curr_cluster_info

    async def get_clusters_read_info(self) -> ClustersInfo:
        raw_response = await self._make_request(
            'DiscoverClusters',
            read_sessions=[ReadSessionParams(all_original={})],
            write_sessions=[WriteSessionParams()],
        )

        parsed_result = DiscoverClustersResult()
        parsed_result.ParseFromString(raw_response)

        result = {}

        # logbroker does NOT disable clusters on read for now, that's why we just merge check results
        self._write_dc_availability(result, parsed_result.write_sessions_clusters)
        self._write_dc_availability(result, parsed_result.read_sessions_clusters)

        return ClustersInfo(result)
