from streamlit_app.common.local_connection import get_connection_parameters
import sys
import path
from snowflake.snowpark import Session, Window
import snowflake.snowpark.table
from snowflake.snowpark.functions import (
    avg,
    col,
    count,
    countDistinct,
    min,
    sum,
    to_date,
)
from typing import List


def filter_relevant_records(
    df: snowflake.snowpark.table.Table, stores: List[str]
) -> snowflake.snowpark.table.Table:
    """
    Function to apply filters to remove unnecessary data
    :param df: Snowpark dataframe
    :param stores: List of Store ids  to filter on
    :return: Snowpark dataframe
    """
    result: snowflake.snowpark.table.Table = df.filter(
        (col("ORDER_TOTAL") > 0)
        & (col("DELETE_REQUESTED") == 0)
        & (col("ORDER_STATUS") == "paid")
        & (col("S_STORE_ID").isin(stores))
    )

    return result


def calculate_simple_metrics(
    df: snowflake.snowpark.table.Table,
) -> snowflake.snowpark.table.Table:
    """
    Function to calculate summary statistics per group
    :param df: Snowflake dataframe with group by
    :return: Snowflake dataframe
    """
    summary_df: snowflake.snowpark.table.Table = df.agg(
        sum("order_items_total").alias("sum_order_total"),
        avg("order_total").alias("avg_order_total"),  # Might include duplicates
        count("order_id").alias("total_orders"),
        countDistinct("store_customer_id").alias("unique_customers"),
        sum("order_items_quantity").alias("total_items_purchased"),
    )
    return summary_df


def calculate_running_totals_territory(
    df: snowflake.snowpark.table.Table,
) -> snowflake.snowpark.table.Table:
    """
    Function to calculate running totals per group
    :param df: Snowflake dataframe with group by
    :return: Snowflake dataframe
    """

    min_date = df.agg(min(to_date(col("order_date_created"))).alias("first_date"))
    new_customers = min_date.groupBy("artist", "territory", "first_date").agg(
        count("store_customer_id").alias("new_customers")
    )
    running_total: snowflake.snowpark.table.Table = new_customers.with_column(
        "RUNNING_TOTAL",
        sum("new_customers").over(
            Window.partition_by([col("artist"), col("territory")]).order_by(col("first_date"))
        ),
    ).select("artist", "territory", "first_date", "running_total")
    return running_total


def create_sales_metrics(sn_session: Session, store_ids: List[str]) -> None:
    """
    Function to create sales metrics data
    :param sn_session: Snowflake session object
    :param store_ids: List of Shopify store ids
    :return: None
    """

    raw_by_artist = "merch_artist_raw_metrics"
    totals_by_artist = "merch_artist_metrics"
    totals_by_artist_daily = "merch_daily_metrics"
    delphi_orders = "CRM_ECOMMERCE_DATA.CONSOLIDATION_DATA.ECOMMERCE_ORDERS"
    delphi_stores = "CRM_ECOMMERCE_DATA.CONSOLIDATION_DATA.ECOMMERCE_STORES"
    delphi_fans = "DELPHI_CRM_DATA.RAW_SALESFORCE_SALES_CLOUD.FAN_C"
    # define general joins
    orders = (
        sn_session.table(delphi_orders)
        .select(
            "store_id",
            "store_customer_id",
            "order_total",
            "order_items_total",
            "order_items_quantity",
            "order_id",
            "delete_requested",
            "order_status",
            "order_date_created",
            "customer_id",
            to_date(col("order_date_created")).alias("order_date"),
        )
        .rename(col("store_id"), "o_store_id")
    )

    stores = (
        sn_session.table(delphi_stores)
        .select("artist", "store_id", "label_id", "territory")
        .rename(col("store_id"), "s_store_id")
    )

    fans = (
        sn_session.table(delphi_fans)
        .select("age_range_c", "gender_c", "customer_id_c")
    )

    orders_join = orders.join(stores, orders.O_STORE_ID == stores.S_STORE_ID)
    join = orders_join.join(fans,  orders_join.CUSTOMER_ID== fans.CUSTOMER_ID_C, how="left")

    query_filter = filter_relevant_records(join, store_ids)

    # save raw data
    query_filter.write.mode("overwrite").save_as_table(raw_by_artist)

    # Calculate simple statistics
    group = query_filter.groupBy(["artist", "territory"])
    agg1 = calculate_simple_metrics(group)

    agg1.write.mode("overwrite").save_as_table(totals_by_artist)

    # Calculate daily sales
    group2 = query_filter.groupBy("artist", "territory", "order_date")
    agg2 = group2.agg(
        sum("order_items_total").alias("sum_order_total"),
        countDistinct("order_id").alias("total_orders"),
    )
    sort2 = agg2.sort(col("territory"), col("order_date").asc())

    # Calculate running totals
    group_running_total = query_filter.groupBy(["artist", "territory", "store_customer_id"])

    results_running_total = calculate_running_totals_territory(group_running_total).sort(
        col("ARTIST"), col("TERRITORY"), col("FIRST_DATE").asc()
    )
    left_join_results = sort2.join(
        results_running_total,
        (sort2.order_date == results_running_total.first_date)
        & (sort2.artist == results_running_total.artist)
        & (sort2.territory == results_running_total.territory),
    ).select(
        sort2.artist.alias("artist"),
        sort2.territory.alias("territory"),
        sort2.order_date,
        sort2.sum_order_total,
        sort2.total_orders,
        results_running_total.running_total,
    )

    left_join_results.write.mode("overwrite").save_as_table(totals_by_artist_daily)


# directory reach
directory = path.Path(__file__).abspath()
# setting path
sys.path.append(directory.parent.parent)


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