from qset_feature_store.repository import models

# todo: refactor code duplicate


class DatasetManager:
    def _format_generic_featureset(self, generic_featureset):
        if generic_featureset['_cls'] == 'DataSet':
            return models.cast_dict(generic_featureset, models.DataSet)
        elif generic_featureset['_cls'] == 'FeatureSet':
            return models.cast_dict(generic_featureset, models.FeatureSet)
        elif generic_featureset['_cls'] == 'DatasetSchema':
            return models.cast_dict(generic_featureset, models.DatasetSchema)

    def get(self, generic_featureset, config=None, params=None):
        generic_featureset = self._format_generic_featureset(generic_featureset)
        if generic_featureset['_cls'] == 'DataSet':
            return self.get_dataset(generic_featureset)
        elif generic_featureset['_cls'] == 'FeatureSet':
            return self.get_from_featureset(generic_featureset, config, params)
        elif generic_featureset['_cls'] == 'DatasetSchema':
            assert params is None
            return self.get_from_dataset_schema(generic_featureset, config)

    def put(self, df, generic_featureset, scope=None):
        generic_featureset = self._format_generic_featureset(generic_featureset)
        if generic_featureset['_cls'] == 'DataSet':
            return self.put_dataset(df, generic_featureset)
        elif generic_featureset['_cls'] == 'DatasetSchema':
            return self.put_to_dataset_schema(df, generic_featureset, scope)

    def get_dataset(self, dataset):
        raise NotImplementedError

    def get_from_featureset(self, featureset, config, params):
        raise NotImplementedError

    def get_from_dataset_schema(self, dataset_schema, config):
        raise NotImplementedError

    def put_dataset(self, df, dataset, **kwargs):
        raise NotImplementedError

    def put_to_dataset_schema(self, df, dataset_schema, config, **kwargs):
        raise NotImplementedError