from typing import Optional
from utils.data_model_utils import generate_collection
from .enrichment import Enrichment
from .enrichment import component_logger


class MissingInputFieldException(Exception):
    pass


class Algo(Enrichment):
    # create segments based on this column
    segment_attribute_name: Optional[str] = None
    segment_name = "Segment for {}"
    run_generate_segments = True

    def _get_unique_attribute_values(self, attribute_name):
        return self.result_df[attribute_name].unique()

    def _populate_collection_table(self):
        """ Generating new collection """
        enrichment_collection_id = generate_collection(
            collection_name=self.enrichment_name,
            collection_source=self.enrichment_name,
            schema_name=self.schema,
            user_id=self.user_id,
            parent_id=self.base_collection_id,
            collection_type="algo",
            log_description="initiate_enrichment",
            status="initiated",
        )
        self.enrichment_collection_id = enrichment_collection_id

    def _get_segment_spec(self):
        if self.run_generate_segments:
            try:
                return {
                    "segment_attribute": self.segment_attribute_name,
                    "segment_name": self.segment_name,
                    "segment_values": self._get_unique_attribute_values(
                        self.segment_attribute_name
                    ).tolist(),
                }
            except KeyError as e:
                component_logger.info(
                    f"{self.enrichment_name} cannot generate segments because {e}"
                )
            except RuntimeError as e:
                component_logger.exception(
                    f"{self.enrichment_name} RuntimeError", exc_info=e
                )
        return {}

    def do_everything(self):
        """
            Add segment creation specification to return values
        :return:
        """
        the_usual = super().do_everything()
        if the_usual:
            the_usual.update(self._get_segment_spec())
        return the_usual

    def _check_collection_type(self):
        """ Simple evaluation of the collection_type """
        # !TODO need to add better evaluation
        if self.collection_type not in ["source", "set", "segment"]:
            """ Making sure that we don't run segments generation if the above logic did not run"""
            self.run_generate_segments = False
            raise RuntimeError("Cannot enrich this type of collection.")

    def _check_available_system_fields(self):
        """ Basic attributes check - will not run the enrichment if source data
            doesn't have at least one available self.source_attribute_od
            """
        ok_to_continue_enrichment = True
        missing_source_attribute_ids = list(
            set(self.source_attribute_ids) - set(self.available_system_fields)
        )
        if len(missing_source_attribute_ids) == len(self.source_attribute_ids):
            ok_to_continue_enrichment = False
            """ Making sure that we don't run segments generation if the above logic did not run"""
            self.run_generate_segments = False
            info_string = f"{self.schema} - Cannot run {self.enrichment_name} - no source attributes available"
            component_logger.info(info_string)
            self._update_collection_status("failed")
            """ We want to return False because we want to use this to exit enrichment process. We don't want to 
            return an exception because this is handled and ok, and we don't want to spam alerts """
            # raise MissingInputFieldException(info_string)
        return ok_to_continue_enrichment