from __future__ import annotations

import typing as t

from sarus_differential_privacy.query import PrivateQuery

from sarus_data_spec.constants import PUBLIC
from sarus_data_spec.sarus_statistics.protobuf.bounds_pb2 import (
    BoundsParameters,
)
from sarus_data_spec.sarus_statistics.protobuf.marginals_pb2 import (
    MarginalsParameters,
)
from sarus_data_spec.sarus_statistics.protobuf.multiplicity_pb2 import (
    MultiplicityParameters,
)
from sarus_data_spec.sarus_statistics.protobuf.size_pb2 import SizeParameters
import sarus_data_spec.typing as st

StatisticsProtobuf = t.Union[
    MultiplicityParameters,
    SizeParameters,
    BoundsParameters,
    MarginalsParameters,
]


class StatisticsParameters(t.Protocol):
    """A python class to describe statistics parameters"""

    def prototype(self) -> t.Type: ...

    def protobuf(self) -> StatisticsProtobuf: ...

    def compute(self, dataset: st.Dataset) -> st.Statistics: ...

    def children(self) -> t.Mapping[str, StatisticsParameters]:
        """Returns the children contained in the type tree structure"""
        ...

    def type(self) -> str: ...

    def private_query(self) -> PrivateQuery: ...

    def random_seed(self) -> int: ...


class BaseStatisticsParameters(StatisticsParameters):
    def type(self) -> str:
        proto_type = self.protobuf().WhichOneof("type")
        if not proto_type:
            raise ValueError("Protobuf does not have a type")
        return proto_type

    def random_seed(self) -> int:
        return self.protobuf().random_seed


def get_is_public(properties: t.Optional[t.Mapping[str, str]]) -> bool:
    """Returns is_public boolean from properties"""
    return bool(properties and properties.get(PUBLIC) == "True")