import copy

import pandas as pd

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

from .external_service_enrichment import ServiceEnrichment


class EnrichKlaviyoStaticData(ServiceEnrichment):
    source_attribute_ids = [111]  # klaviyoPersonId
    result_attribute_ids = [
        21,  # userLocationLon
        20,  # userLocationLat
        15,  # userAddress
        # 19,   # userAddress2
        # 16,   # userArea
        4,  # userLastName
        6,  # userLastName
        12,  # userCountry
        18,  # userZip
        2,  # userFirstName
        13,  # userCity
    ]
    data_gathering_rule = "strict"
    get_previously_enriched_data = False

    final_columns = {
        "$longitude": "userLocationLon",
        "$latitude": "userLocationLat",
        "$address1": "userAddress",
        # '$address2': 'userAddress2',
        "$region": "userState",
        "$last_name": "userLastName",
        "$phone_number": "userPhone",
        "$country": "userCountry",
        "$zip": "userZip",
        "$first_name": "userFirstName",
        "$city": "userCity",
        "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_static_data_async(person_id=person_id)
            )

        for task in async_response_list:
            response = task.wait_for_result()
            data = response["response"]
            person_id = response["person_id"]
            if "data" in data:
                data["data"]["klaviyo_person_id"] = person_id
                response_list.append(data["data"])

        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", "row_id"], inplace=True)

        self.result_df = result_df

    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.available_data["source"][["fan_id", "row_id"]].merge(
            self.result_df, how="left", on="fan_id"
        )

        # 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