import pandas as pd

from service.utils.aws_connectors import run_query

from .enrichment import Enrichment


class ServiceEnrichment(Enrichment):
    """General Class for all the Service Enrichments
    requires management_schema to be provided in order to obtain service credentials"""

    def __init__(
        self, collection_id, user_id: str, schema, default_value, management_schema
    ):
        super().__init__(
            collection_id=collection_id,
            schema=schema,
            user_id=user_id,
            default_value=default_value,
        )
        self.management_schema = management_schema


class EventServiceEnrichment(ServiceEnrichment):
    """
    Extending the Enrichment class in order to support custom logic of importing data from External Services like Klaviyo
    Adding the missing rows in to source collection to match the output of the enrichment.
    This happens when the source collection is a list of unique fans and enrichment can retrieve multiple events per each fan.
    """

    def __init__(
        self, collection_id, user_id: str, schema, default_value, management_schema
    ):
        super().__init__(
            collection_id=collection_id,
            schema=schema,
            user_id=user_id,
            default_value=default_value,
            management_schema=management_schema,
        )
        self.update_source_df = pd.DataFrame({})

    def _update_source_collection(self):
        """Updating the source collection data:
        Adding extra rows to match the self.update_source_df
        so that source collection would have exactly the same rows as the enrichment results
        that can have multiple events per row
        """
        if len(self.update_source_df) > 0:
            params = {
                "source_collection_id": self.base_collection_id,
                "enrichment_collection_id": self.enrichment_collection_id,
            }

            sql = f"""
            INSERT INTO {self.schema}.fan_attribute (
                WITH source AS (
                    SELECT *
                    FROM {self.schema}.fan_attribute
                    WHERE collection_id = %(source_collection_id)s
                ),
                     duplicate AS (
                         SELECT DISTINCT fan_id, row_id
                         FROM {self.schema}.fan_attribute
                         WHERE collection_id = %(enrichment_collection_id)s
                ),
                     base AS (
                         SELECT d.*
                         FROM duplicate d
                         LEFT JOIN source s
                            ON d.fan_id = s.fan_id
                           AND d.row_id = s.row_id
                         WHERE s.fan_id IS NULL
                )
                SELECT s.fan_id,
                       s.attribute_id,
                       s.value,
                       b.row_id,
                       s.collection_id,
                       current_timestamp AS timestamp
                FROM base b
                         LEFT JOIN source s
                                   ON b.fan_id = s.fan_id
            );
            """
            run_query(sql, params)
        else:
            pass

    def _finish(self):
        """Updating source collection data once the enrichment finishes"""
        super()._finish()
        self._update_source_collection()