import copy

import pandas as pd

from service.tasks.api_service_handler.klaviyo.service_klaviyo import KlaviyoServiceCore
from service.tasks.api_service_handler.utils import flatten_json
from service.tasks.base_attributes_generation import (
    get_fan_attribute_values,
    load_new_attributes_to_db,
)

from .external_service_enrichment import EventServiceEnrichment


class EnrichKlaviyoEventsData(EventServiceEnrichment):
    source_attribute_ids = [111]  # klaviyoPersonId

    #  This Enrichment producess more fields than states here.
    #  In addition to 2 pre defined system fields we are storing all non system fields for further evaluation
    result_attribute_ids = [22, 29]  # eventName  # eventPurchaseDate
    data_gathering_rule = "strict"
    get_previously_enriched_data = False

    final_columns = {"klaviyo_person_id": "klaviyoPersonId"}

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

    # @timeit
    def _run(self):
        """Main enrichment method - rewriting parent._run() method for enrichment to work"""
        source_df = copy.deepcopy(self.available_data["source"])

        """ Making sure that we have one row per fan_id"""
        if len(source_df) > len(source_df.fan_id.unique()):
            source_df = (
                source_df.reset_index(drop=True)
                .fillna("")
                .groupby("fan_id")
                .max()
                .reset_index()
            )

        ids_list = source_df["111"].unique()
        klaviyo = KlaviyoServiceCore(schema=self.management_schema)
        response_list = []
        async_response_list = []

        for person_id in ids_list:
            async_response_list.append(
                klaviyo.get_person_events_data_async(person_id=person_id)
            )

        for task in async_response_list:
            response = task.wait_for_result()
            for event in response:
                try:
                    """Removing personal details as those will be obtained
                    through other enrichment"""
                    del event["person"]
                except Exception:
                    pass
                event = flatten_json(event)
                response_list.append(event)
        if len(response_list) > 0:
            result_df = pd.DataFrame(response_list)
            result_df.rename(columns=self.final_columns, inplace=True)

            result_df = result_df.merge(
                source_df, how="left", right_on="111", left_on="klaviyoPersonId"
            )
            result_df.drop(columns=["111", "klaviyoPersonId"], inplace=True)

            """Manipulating the data to ensure that we match row_ids for the source data
               + create extra rows in source collection"""
            """Getting row count per each fan_id"""
            result_df["rank"] = (
                result_df.sort_values(["row_id"], ascending=[True])
                .groupby(["fan_id"])
                .cumcount()
            )

            # these are events that can go without any changes
            unchanged_df = result_df[result_df["rank"] == 0]

            # these are events that have to change their row_ids and adjust source collection (add rows to match the enrichment)
            df_to_change = result_df[result_df["rank"] > 0]
            df_to_change.reset_index(drop=True, inplace=True)
            df_to_change["row_id"] = df_to_change.index + df_to_change.row_id.max() + 1
            self.update_source_df = copy.deepcopy(df_to_change[["fan_id", "row_id"]])

            result_df = pd.concat([unchanged_df, df_to_change], ignore_index=True)
            result_df.drop(columns=["rank"], inplace=True)
            self.result_df = result_df
        else:
            self.result_df = pd.DataFrame({})

    def _generate_fan_attributes(self):
        """Getting a list of attributes that are present in this enrichment
        Assigning none system attribute_ids to the extra fields.
        """
        condition = self.system_fields["a_id"].isin(self.result_attribute_ids)
        attributes_df = self.system_fields[condition][["a_system_name", "a_id"]]
        attributes_df.rename(
            columns={"a_system_name": "name", "a_id": "attribute_id"}, inplace=True
        )

        # We make sure we retain the original row_id
        new_result_df = self.result_df

        # If we don't set index, the df default index will be used (and thus wrong row_ids inserted)
        new_result_df.set_index("row_id", inplace=True)

        # We need to add new non_system_fields to attribute table and return their IDs to get_fan_attribute_values
        attributes_df, current_attributes_df = load_new_attributes_to_db(
            schema=self.schema,
            source_df=self.result_df,
            sys_fields=list(self.final_columns.values()),
        )

        fan_attributes_df = get_fan_attribute_values(
            source_df=new_result_df,
            attributes_df=current_attributes_df,
            collection_id=self.enrichment_collection_id,
        )

        # Now when we know all the attribute_ids we need to update this field for populating data in `self.finish()`
        self.result_attribute_ids = list(
            current_attributes_df[current_attributes_df["name"] != "fan_id"][
                "attribute_id"
            ]
        )

        self.fan_attribute = fan_attributes_df

    def _populate_attribute_table(self):
        """We don't need to populate attribute_table separately
        it is handled by load_new_attributes_to_db() in _generate_fan_attributes()
        """
        pass