# Columns which expected to be in the output table
# They used by TADAS_METRICS DBT model
import logging
import os

import pandas as pd

from tadas.platform import metrics as metrics_utils
from tadas.domain.constants import REPORT_FINAL_TABLES
from tadas.platform import context as contexts
from tadas.snowflake import client as snowflake_utils
from tadas.platform import config

logger = logging.getLogger(__name__)

EXPECTED_COLUMNS = [
    'isrc_cd',
    'geo_country',
    'report_date',
    'consecutive_trend',
    'tadas_30day_score',
    'spotify_collection_flag',
    'spotify_collection_lift',
    'spotify_lean_forward_flag',
    'spotify_lean_forward_lift',
    'spotify_search_flag',
    'spotify_search_lift',
    'spotify_all_flag',
    'spotify_all_lift',
    'apple_lean_forward_flag',
    'apple_lean_forward_lift',
    'apple_search_flag',
    'apple_search_lift',
    'apple_all_flag',
    'apple_all_lift',
    'shazam_flag',
    'shazam_lift',
    'tiktok_creations_country_flag',
    'tiktok_creations_country_lift',
    'tiktok_creations_global_flag',
    'tiktok_creations_global_lift',
    'tiktok_views_country_flag',
    'tiktok_views_country_lift',
    'tiktok_views_global_flag',
    'tiktok_views_global_lift',
    'tiktok_likes_country_flag',
    'tiktok_likes_country_lift',
    'tiktok_likes_global_flag',
    'tiktok_likes_global_lift',
    'insta_reels_views_country_flag',
    'insta_reels_views_country_lift',
    'insta_reels_views_global_flag',
    'insta_reels_views_global_lift',
    'insta_reels_creations_country_flag',
    'insta_reels_creations_country_lift',
    'insta_reels_creations_global_flag',
    'insta_reels_creations_global_lift',
]

TRENDING_FLAGS_COLUMNS = [v for v in EXPECTED_COLUMNS if v.endswith('_flag')]

def get_source_report_date(source_table):
    """Return the single report_date stored in the given output table."""
    query = f"SELECT DISTINCT report_date FROM {source_table}"
    with snowflake_utils.snowflake_connection() as conn:
        result = snowflake_utils.query_snowflake_fetchall(connection=conn, query=query)
    dates = [row[0] for row in result]
    assert len(dates) == 1, f"Expected exactly one report_date in {source_table}, got {dates}"
    return dates[0]


def copy_from_model_outputs_to_target(source_table, target_table, force=False):
    """
    Copy from model outputs to target table. It doesn't copy if:
    * the source_table doesn't have trending data
    * number of rows in source table equals to target table

    :param source_table:
    :param target_table:
    :param force:
    :return: True if the data was copied, else False
    """

    common_source_params = {
        'consecutive_trend_upper_limit': config.get('TRENDING_DAYS_UPPER_LIMIT'),
    }
    common_source_filters = f"""
        consecutive_trend between 0 and %(consecutive_trend_upper_limit)s
        and tadas_30day_score > 0
    """

    query_rows_in_source = f"""
    select *
    from {source_table}
    where {common_source_filters}
    """
    with snowflake_utils.snowflake_connection() as connection:
        # get counts from the source table
        source_df_filtered = snowflake_utils.query_snowflake_to_df(
            connection=connection,
            query=query_rows_in_source,
            params=common_source_params,
        )
        total_rows_in_source = len(source_df_filtered)
        logger.info(f"Rows in source table {source_table}: {total_rows_in_source}")
        if not total_rows_in_source:
            return False

        if not snowflake_utils.is_table_exists(target_table, conn=connection):
            total_rows_in_target = 0
        else:
            # compare with counts from the target table
            query_rows_in_target = f"""
            select count(*) as total_rows
            from {target_table}
            """
            result = snowflake_utils.query_snowflake_fetchall(
                connection=connection,
                query=query_rows_in_target,
                params=common_source_params,
            )
            total_rows_in_target = result[0][0] if result else 0

        logger.info(f"Rows in target table {target_table}: {total_rows_in_target}")
        if total_rows_in_target == total_rows_in_source and not force:
            logger.info("Already in sync. Exiting.")
            return False

        # re-create the target if needed
        query_sync_to_target = f"""
        create or replace transient table {target_table} as
        select {','.join(EXPECTED_COLUMNS)}
        from {source_table}
        where 
             {common_source_filters}
        """
        snowflake_utils.query_snowflake_fetchall(
            connection=connection,
            query=query_sync_to_target,
            params=common_source_params,
        )

        # send metrics
        metrics = metrics_utils.Metrics(
            category='output_totals',
            context=contexts.load_context(),
        )
        metrics.add_metric("trending_tracks_count", total_rows_in_source)
        metrics.send()

        # send metrics breakdown by country
        # aggregate df by country
        df_aggregated = source_df_filtered.groupby(['geo_country']).agg({'isrc_cd': 'count'}).reset_index()
        # iterate over countries and send metrics
        metrics_per_country = metrics_utils.Metrics(
            category='output_per_country',
            context=contexts.load_context(),
        )
        for row in df_aggregated.itertuples():
            metrics_per_country.add_metric('country', row.geo_country)
            metrics_per_country.add_metric('trending_tracks_count', row.isrc_cd)
            metrics_per_country.send()

        return True


def sync_historical(source_table, target_table, model_version):
    """
    Sync current model output to the historical table.

    1. Delete from target_table all entries matching current source_table
       (by report_date and geo_country).
    2. Insert selected columns from source_table into target_table with
       the given model_version.
    """
    with snowflake_utils.snowflake_connection() as connection:
        # Delete existing entries for the same report_date + country + model_version
        delete_query = f"""
        DELETE FROM {target_table}
        WHERE model_version = %(model_version)s
          AND (report_date, geo_country) IN (
              SELECT DISTINCT report_date, geo_country
              FROM {source_table}
          )
        """
        snowflake_utils.execute(
            connection, delete_query, params={'model_version': model_version}
        )
        logger.info(f"Deleted overlapping entries from {target_table}")

        # Copy data from source to target
        insert_query = f"""
        INSERT INTO {target_table}
            (report_date, isrc_cd, geo_country, pfn_geo,
             tadas_30day_score, consecutive_trend, model_version)
        SELECT
            report_date,
            isrc_cd,
            geo_country,
            pfn_geo,
            tadas_30day_score,
            consecutive_trend,
            %(model_version)s
        FROM {source_table}
        """
        snowflake_utils.execute(
            connection, insert_query, params={'model_version': model_version}
        )
        logger.info(
            f"Copied data from {source_table} to {target_table} "
            f"with model_version={model_version}"
        )


def load_final_tables_as_df():
    result = {}
    with snowflake_utils.snowflake_connection() as conn:
        for report, table_name in REPORT_FINAL_TABLES.items():
            df = snowflake_utils.query_snowflake_to_df(conn, f'SELECT * FROM {table_name}', use_cache=False)
            result[report] = df
    return result


def save_combined_df(combined_df: pd.DataFrame, table_name: str):
    combined_df["report_date"] = pd.to_datetime(combined_df["report_date"], format="%Y-%m-%d").dt.date
    snowflake_utils.saveto_snowflake(
        df=combined_df,
        myschema=os.getenv('SNOWFLAKE_SCHEMA'),
        table=table_name,
        mode='replace',
        fix_column_names=True,
    )