from abc import abstractmethod
from datetime import datetime
from typing import TYPE_CHECKING, Dict, List, Optional, Type, cast

import pandas as pd
import pyarrow
from google.protobuf.json_format import MessageToJson

from feast.data_source import DataSource
from feast.dqm.profilers.profiler import Profile, Profiler
from feast.protos.feast.core.SavedDataset_pb2 import SavedDataset as SavedDatasetProto
from feast.protos.feast.core.SavedDataset_pb2 import SavedDatasetMeta, SavedDatasetSpec
from feast.protos.feast.core.SavedDataset_pb2 import (
    SavedDatasetStorage as SavedDatasetStorageProto,
)

if TYPE_CHECKING:
    from feast.infra.offline_stores.offline_store import RetrievalJob


class _StorageRegistry(type):
    classes_by_proto_attr_name: Dict[str, Type["SavedDatasetStorage"]] = {}

    def __new__(cls, name, bases, dct):
        kls = type.__new__(cls, name, bases, dct)
        if dct.get("_proto_attr_name"):
            cls.classes_by_proto_attr_name[dct["_proto_attr_name"]] = kls
        return kls


class SavedDatasetStorage(metaclass=_StorageRegistry):
    _proto_attr_name: str

    @staticmethod
    def from_proto(storage_proto: SavedDatasetStorageProto) -> "SavedDatasetStorage":
        proto_attr_name = cast(str, storage_proto.WhichOneof("kind"))
        return _StorageRegistry.classes_by_proto_attr_name[proto_attr_name].from_proto(
            storage_proto
        )

    @abstractmethod
    def to_proto(self) -> SavedDatasetStorageProto:
        ...

    @abstractmethod
    def to_data_source(self) -> DataSource:
        ...


class SavedDataset:
    name: str
    features: List[str]
    join_keys: List[str]
    full_feature_names: bool
    storage: SavedDatasetStorage
    tags: Dict[str, str]
    feature_service_name: Optional[str] = None

    created_timestamp: Optional[datetime] = None
    last_updated_timestamp: Optional[datetime] = None

    min_event_timestamp: Optional[datetime] = None
    max_event_timestamp: Optional[datetime] = None

    _retrieval_job: Optional["RetrievalJob"] = None

    def __init__(
        self,
        name: str,
        features: List[str],
        join_keys: List[str],
        storage: SavedDatasetStorage,
        full_feature_names: bool = False,
        tags: Optional[Dict[str, str]] = None,
        feature_service_name: Optional[str] = None,
    ):
        self.name = name
        self.features = features
        self.join_keys = join_keys
        self.storage = storage
        self.full_feature_names = full_feature_names
        self.tags = tags or {}
        self.feature_service_name = feature_service_name

        self._retrieval_job = None

    def __repr__(self):
        items = (f"{k} = {v}" for k, v in self.__dict__.items())
        return f"<{self.__class__.__name__}({', '.join(items)})>"

    def __str__(self):
        return str(MessageToJson(self.to_proto()))

    def __hash__(self):
        return hash((id(self), self.name))

    def __eq__(self, other):
        if not isinstance(other, SavedDataset):
            raise TypeError(
                "Comparisons should only involve FeatureService class objects."
            )
        if self.name != other.name:
            return False

        if sorted(self.features) != sorted(other.features):
            return False

        return True

    @staticmethod
    def from_proto(saved_dataset_proto: SavedDatasetProto):
        """
        Converts a SavedDatasetProto to a SavedDataset object.

        Args:
            saved_dataset_proto: A protobuf representation of a SavedDataset.
        """
        ds = SavedDataset(
            name=saved_dataset_proto.spec.name,
            features=list(saved_dataset_proto.spec.features),
            join_keys=list(saved_dataset_proto.spec.join_keys),
            full_feature_names=saved_dataset_proto.spec.full_feature_names,
            storage=SavedDatasetStorage.from_proto(saved_dataset_proto.spec.storage),
            tags=dict(saved_dataset_proto.spec.tags.items()),
        )

        if saved_dataset_proto.spec.feature_service_name:
            ds.feature_service_name = saved_dataset_proto.spec.feature_service_name

        if saved_dataset_proto.meta.HasField("created_timestamp"):
            ds.created_timestamp = (
                saved_dataset_proto.meta.created_timestamp.ToDatetime()
            )
        if saved_dataset_proto.meta.HasField("last_updated_timestamp"):
            ds.last_updated_timestamp = (
                saved_dataset_proto.meta.last_updated_timestamp.ToDatetime()
            )
        if saved_dataset_proto.meta.HasField("min_event_timestamp"):
            ds.min_event_timestamp = (
                saved_dataset_proto.meta.min_event_timestamp.ToDatetime()
            )
        if saved_dataset_proto.meta.HasField("max_event_timestamp"):
            ds.max_event_timestamp = (
                saved_dataset_proto.meta.max_event_timestamp.ToDatetime()
            )

        return ds

    def to_proto(self) -> SavedDatasetProto:
        """
        Converts a SavedDataset to its protobuf representation.

        Returns:
            A SavedDatasetProto protobuf.
        """
        meta = SavedDatasetMeta()
        if self.created_timestamp:
            meta.created_timestamp.FromDatetime(self.created_timestamp)
        if self.min_event_timestamp:
            meta.min_event_timestamp.FromDatetime(self.min_event_timestamp)
        if self.max_event_timestamp:
            meta.max_event_timestamp.FromDatetime(self.max_event_timestamp)

        spec = SavedDatasetSpec(
            name=self.name,
            features=self.features,
            join_keys=self.join_keys,
            full_feature_names=self.full_feature_names,
            storage=self.storage.to_proto(),
            tags=self.tags,
        )
        if self.feature_service_name:
            spec.feature_service_name = self.feature_service_name

        feature_service_proto = SavedDatasetProto(spec=spec, meta=meta)
        return feature_service_proto

    def with_retrieval_job(self, retrieval_job: "RetrievalJob") -> "SavedDataset":
        self._retrieval_job = retrieval_job
        return self

    def to_df(self) -> pd.DataFrame:
        if not self._retrieval_job:
            raise RuntimeError(
                "To load this dataset use FeatureStore.get_saved_dataset() "
                "instead of instantiating it directly."
            )

        return self._retrieval_job.to_df()

    def to_arrow(self) -> pyarrow.Table:
        if not self._retrieval_job:
            raise RuntimeError(
                "To load this dataset use FeatureStore.get_saved_dataset() "
                "instead of instantiating it directly."
            )

        return self._retrieval_job.to_arrow()

    def as_reference(self, profiler: "Profiler") -> "ValidationReference":
        return ValidationReference(profiler=profiler, dataset=self)

    def get_profile(self, profiler: Profiler) -> Profile:
        return profiler.analyze_dataset(self.to_df())


class ValidationReference:
    dataset: SavedDataset
    profiler: Profiler

    def __init__(self, dataset: SavedDataset, profiler: Profiler):
        self.dataset = dataset
        self.profiler = profiler

    @property
    def profile(self) -> Profile:
        return self.profiler.analyze_dataset(self.dataset.to_df())
