import pandas as pd
from snowflake.snowpark import Session
from snowflake.snowpark.functions import lower, col, sha2, when, lit, countDistinct, count, sum

from streamlit_app.common.local_connection import get_connection_parameters


def create_fan_insights(sn_session: Session) -> bool:
    """
    Function to create fan insights data based on purchases and Fansifter data
    :param sn_session: Snowflake session object
    :return: None
    """

    output_table = "merch_artist_fansifter_percentage"
    output_table_segments = "merch_artist_segment_percentage"
    delphi_orders = "CRM_ECOMMERCE_DATA.CONSOLIDATION_DATA.ECOMMERCE_ORDERS"
    delphi_stores = "CRM_ECOMMERCE_DATA.CONSOLIDATION_DATA.ECOMMERCE_STORES"
    crm_fans_table = "DELPHI_CRM_DATA.RAW_SALESFORCE_SALES_CLOUD.FAN_C"
    fansifter_fans = "FANSIFTER_APP_REPORTING.PROD.FAN_ARTIST_DBT"
    fansifter_segments = "FANSIFTER_APP_REPORTING.PROD.GLOBAL_FAN_SEGMENT_DBT"

    # mapping data
    store_list = ['S00030', 'S00088', 'S00113', 'S00115', 'S00121', 'S00122', 'S00123']
    artist_list = ['4498fe09-0150-4e54-a201-0d52ffc56bcd', '4498fe09-0150-4e54-a201-0d52ffc56bcd',
                   'bfbe7cba-93f1-4c0b-96a2-819397540c04',
                   '06183d1d-7c5b-444b-87ef-6a9231378f0b', 'bfbe7cba-93f1-4c0b-96a2-819397540c04',
                   'bfbe7cba-93f1-4c0b-96a2-819397540c04',
                   'bfbe7cba-93f1-4c0b-96a2-819397540c04']
    artist_names = ['Rex Orange County', 'Rex Orange County', 'LISA', 'SZA', 'LISA', 'LISA', 'LISA']
    mappings = {
        'GP_ID': artist_list,
        'STORE_ID': store_list,
        'ARTIST_NAME': artist_names
    }

    sp_mappings = sn_session.createDataFrame(pd.DataFrame(mappings))

    orders = (
        sn_session.table(delphi_orders)
        .select(
            "store_id",
            "customer_id",
        ).rename(col("store_id"), "o_store_id").filter((col("O_STORE_ID").isin(store_list)))
    )
    stores = (
        sn_session.table(delphi_stores)
        .select(
            "store_id",
            "territory",
        ).rename(col("store_id"), "s_store_id")
    )
    fans = (
        sn_session.table(fansifter_fans)
        .select(
            "fan_id",
            "global_participant_id",
        )
    ).filter((col("GLOBAL_PARTICIPANT_ID").isin(artist_list)))
    crm_fans  = (
        sn_session.table(crm_fans_table)
        .select(
            "email_c",
            "customer_id_c",
        )
    ).with_column(
        "sha_fan_id", sha2(lower(col("email_c")), 256)
    )

    joined_df = (
        sp_mappings.join(orders, orders.o_store_id == sp_mappings.store_id, how="inner").
        join(stores, orders.o_store_id == stores.s_store_id, how="inner").
        join(crm_fans, orders.customer_id==crm_fans.customer_id_c, how="left").
        join(fans, fans.fan_id==crm_fans.sha_fan_id, how="left")
    ).with_column(
        "FANSIFTER",  when(col("FAN_ID").isNull(), lit(0))
             .otherwise(lit(1))
    ).cache_result()

    aggregation_pre = (joined_df.groupBy(["ARTIST_NAME", "TERRITORY", "FANSIFTER"]).agg(
        countDistinct("CUSTOMER_ID").alias("unique_customers"),
        count("CUSTOMER_ID").alias("total_customers")
    )).cache_result()

    total_unique_per_group = aggregation_pre.groupBy("ARTIST_NAME", "TERRITORY").agg(
        sum("unique_customers").alias("total_unique_customers")
    ).cache_result()

    aggregation = aggregation_pre.join(total_unique_per_group, on=["ARTIST_NAME", "TERRITORY"]).withColumn(
        "percentage", col("unique_customers") / col("total_unique_customers")
    )

    aggregation.write.mode("overwrite").save_as_table(output_table)

    segments = (
        sn_session.table(fansifter_segments)
        .select(
            "fan_id",
            "global_participant_id",
            "segment_name"
        )
    )

    joined_segments_df = (
        sp_mappings.join(orders, orders.o_store_id == sp_mappings.store_id, how="inner").
        join(stores, orders.o_store_id == stores.s_store_id, how="inner").
        join(crm_fans, crm_fans.customer_id_c == orders.customer_id, how="inner").
        join(segments, ((segments.fan_id == crm_fans.sha_fan_id) & (sp_mappings.gp_id == segments.global_participant_id)), how="inner")
    ).cache_result()

    aggregation_segments_pre = joined_segments_df.groupBy(["ARTIST_NAME", "TERRITORY", "SEGMENT_NAME"]).agg(
        countDistinct("CUSTOMER_ID").alias("unique_customers"),
        count("CUSTOMER_ID").alias("total_customers")
    ).cache_result()

    total_unique_segments_per_group = aggregation_segments_pre.groupBy("ARTIST_NAME", "TERRITORY").agg(
        sum("unique_customers").alias("total_unique_customers")
    ).cache_result()

    aggregation_segments = total_unique_segments_per_group.join(aggregation_segments_pre, on=["ARTIST_NAME", "TERRITORY"]).withColumn(
        "percentage", col("unique_customers") / col("total_unique_customers")
    )

    aggregation_segments.write.mode("overwrite").save_as_table(output_table_segments)
    return True


# connection_params = get_connection_parameters("sme_merch")
# session = Session.builder.configs(connection_params).create()
#
# session.sproc.register(
#     create_fan_insights,
#     name="fan_insights_creation",
#     stage_location="@STORED_PROCEDURES",
#     is_permanent=True,
#     execute_as="caller",
#     packages=[
#         "snowflake-snowpark-python==1.28.0",
#     ],
#     replace=True,
# )
#
# session.close()