from dataclasses import dataclass
from typing import List, Optional, Iterable

from ydb.public.api.grpc.draft.ydb_persqueue_v1_pb2_grpc import PersQueueServiceStub
from ydb.public.api.protos.ydb_persqueue_v1_pb2 import Path, ReadInfoRequest, ReadInfoResult

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


@dataclass(frozen=True)
class LogbrokerPartitionInfo:
    end_offset: int
    read_offset: int

    commit_time_lag_ms: int
    read_time_lag_ms: int

    read_session_id: Optional[str] = None

    @classmethod
    def parse(cls, partition) -> 'LogbrokerPartitionInfo':
        return LogbrokerPartitionInfo(
            partition.end_offset,
            partition.read_offset,
            partition.commit_time_lag_ms,
            partition.read_time_lag_ms,
            read_session_id=partition.session_id
        )


@dataclass(frozen=True)
class LogbrokerTopicInfo:
    topic: str
    cluster: str
    partitions: List[LogbrokerPartitionInfo]

    @classmethod
    def parse(cls, topic) -> 'LogbrokerTopicInfo':
        return LogbrokerTopicInfo(
            topic.topic.path,
            topic.cluster,
            [LogbrokerPartitionInfo.parse(partition) for partition in topic.partitions]
        )

    def has_read_session(self) -> bool:
        return any([partition.read_session_id for partition in self.partitions])


class ReadInfoClient(LogbrokerClient):
    _PROTO_MAPPING = {
        'GetReadSessionsInfo': ReadInfoRequest
    }

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

    async def get_read_info(
        self,
        topic_paths: Iterable[str],
        consumer_path: str,
        only_original: bool
    ) -> List[LogbrokerTopicInfo]:
        raw_response = await self._make_request(
            'GetReadSessionsInfo',
            metadata=[('x-ydb-auth-ticket', self._token)],
            topics=map(lambda p: Path(path=p), topic_paths),
            consumer=Path(path=consumer_path),
            get_only_original=only_original
        )

        parsed_result = ReadInfoResult()
        parsed_result.ParseFromString(raw_response)
        return [LogbrokerTopicInfo.parse(topic_info) for topic_info in parsed_result.topics]
