import logging

from service.async_task_manager import io_task
from service.tasks.filtering import create_segment
from service.tasks.filtering.filters import get_filter_conf
from service.utils.data_model_utils import generate_collection

from .enrichment import Enrichment

logger = logging.getLogger(__name__)


class MissingInputFieldException(Exception):
    pass


class Algo(Enrichment):
    # create segments based on this column
    segment_attribute_name = 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 _generate_segments(self):
        """
        Taking unique values of the attribute [segment_attribute_name]
        asking filter configuration for [segment_attribute_name]
        for each unique value asyncronously make segment
        wait for all segments to be created.
        :return:
        """
        if self.run_generate_segments:
            try:
                unique_values = self._get_unique_attribute_values(
                    self.segment_attribute_name
                )
                filter_conf = get_filter_conf(
                    attribute_name=self.segment_attribute_name
                )
                segmenting = [
                    self._create_segment_async(filter_conf["fid"], val)
                    for val in unique_values
                ]
                for seg in segmenting:
                    seg.wait_for_result()
            except KeyError as e:
                logger.info(
                    f"{self.enrichment_name} cannot generate segments because {e}"
                )
            except RuntimeError as e:
                logger.exception(f"{self.enrichment_name} RuntimeError", exc_info=e)
        else:
            logger.info(
                f"Skipping {self.enrichment_name} segments generation as requested by enrichment"
            )

    @io_task
    def _create_segment_async(self, filter_id, filter_value, **kwargs):
        create_segment(
            self.schema,
            self.user_id,
            self.segment_name.format(filter_value),
            self.base_collection_id,
            filter_id,
            [filter_value],
        )

    def do_everything(self):
        super().do_everything()
        self._generate_segments()

    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"
            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