from datetime import datetime, timedelta
import warnings
from pathlib import Path

import dagster as dg
from dagster_snowflake import SnowflakeResource

from src.utils.snowflake.constants import TIME_WINDOW_FRESHNESS_POLICY_WARN_1H_FAIL_2H, PartitionExpr, Warehouse, SnowflakeDB, Team, SnowflakeSchema, Group
from src.utils.snowflake.logger import log_query
from src.utils.snowflake.query import load_query

warnings.filterwarnings("ignore", category=dg.BetaWarning)

# Get directory of this file for relative SQL file loading
ASSET_DIR = Path(__file__).parent

APP_AUDIO_ACTIONS_START_DATE = datetime.strptime('2024-06-01', '%Y-%m-%d')
APP_AUDIO_ACTIONS_TABLE_NAME = "APP_AUDIO_ACTIONS"
WAREHOUSE = Warehouse.WEB_EVENT_X_SMALL.value
SNOWFLAKE_DB = SnowflakeDB.SUNO_PROD.value
SNOWFLAKE_SCHEMA = SnowflakeSchema.PROD.value


@dg.asset(
    name="app_audio_actions",
    description="Legacy iOS audio player actions tracking audio playback events.",
    group_name=Group.SEGMENT_EVENTS.value,
    partitions_def=dg.HourlyPartitionsDefinition(start_date=APP_AUDIO_ACTIONS_START_DATE, end_offset=0),
    backfill_policy=dg.BackfillPolicy.multi_run(max_partitions_per_run=24*7),
    owners=[Team.CORE_POD.value],
    metadata={
        "database": SNOWFLAKE_DB,
        "schema": SNOWFLAKE_SCHEMA,
        "table_name": APP_AUDIO_ACTIONS_TABLE_NAME,
        "data_start_date": APP_AUDIO_ACTIONS_START_DATE.strftime("%Y-%m-%d"),
        "cluster_by": "[p_date, p_hour]",
        "partition_expr": PartitionExpr.HOURLY.value,
        "transient": True,
        "sla_minutes": 15,
    },
    freshness_policy=TIME_WINDOW_FRESHNESS_POLICY_WARN_1H_FAIL_2H,
)
def app_audio_actions(context: dg.AssetExecutionContext, snowflake: SnowflakeResource) -> dg.MaterializeResult:
    run_id = context.run.run_id
    logger = dg.get_dagster_logger()

    # Get partition time window for processing
    partition_start = context.partition_time_window.start
    partition_end = context.partition_time_window.end
    is_multi_partition_range = context.has_partition_key_range

    fetch_params = {
        "partition_start_date": partition_start.strftime("%Y-%m-%d"),
        "partition_end_date": partition_end.strftime("%Y-%m-%d"),
        "partition_start_hour": partition_start.hour,
        "partition_end_hour": partition_end.hour,
        "app_audio_actions_table_name": APP_AUDIO_ACTIONS_TABLE_NAME,
    }

    logger.info(f"Processing app_audio_actions for partition: {partition_start} to {partition_end}")
    logger.info(f"Fetch params: {fetch_params}")

    with snowflake.get_connection() as conn:
        cursor = conn.cursor()
        logger.info(f"Using warehouse {WAREHOUSE}")
        warehouse_query = load_query("src/utils/snowflake/queries/use_warehouse.sql", params={"warehouse": WAREHOUSE})
        log_query(logger, warehouse_query)
        cursor.execute(warehouse_query)

        # 1. Delete existing data for the partition window
        logger.info(f"Deleting existing data from {APP_AUDIO_ACTIONS_TABLE_NAME} for partition window {partition_start} to {partition_end}.")
        delete_query = load_query("src/utils/snowflake/queries/delete_hourly_partitions.sql", params={**fetch_params, "delete_partition_table_name": APP_AUDIO_ACTIONS_TABLE_NAME})
        log_query(logger, delete_query)
        cursor.execute(delete_query)

        # 2. Insert new data hour by hour using insert.sql
        total_rows_inserted = 0
        current_hour = partition_start

        while current_hour < partition_end:
            # Generate file path and geo table name for this hour
            # S3 path format: @SUNO_S3_EVENTS/App-Audio-Player-Action/0/{YEAR}/{MONTH}/{DAY}/{HOUR}
            file_path = f"@SUNO_S3_EVENTS/App-Audio-Action-Event/0/{current_hour.year}/{current_hour.month}/{current_hour.day}/{current_hour.hour}"
            geo_table_name = f"GEO_INFO_TEMP_APP_AUDIO_{current_hour.year}_{current_hour.month}_{current_hour.day}_{current_hour.hour}"

            logger.info(f"Processing hour: {current_hour.strftime('%Y-%m-%d %H:00:00')}")
            logger.info(f"File path: {file_path}")
            logger.info(f"Geo table: {geo_table_name}")

            # Step 3a: Insert data from S3 using insert.sql
            insert_params = {
                "app_audio_actions_table_name": APP_AUDIO_ACTIONS_TABLE_NAME,
                "file_path": file_path,
                "partition_date": current_hour.strftime("%Y-%m-%d"),
                "partition_hour": current_hour.hour,
            }
            insert_query = load_query(ASSET_DIR / "insert.sql", params=insert_params)
            log_query(logger, insert_query)
            cursor.execute(insert_query)

            rows_for_hour = cursor.rowcount
            total_rows_inserted += rows_for_hour
            logger.info(f"Inserted {rows_for_hour} rows for hour {current_hour.strftime('%Y-%m-%d %H:00:00')}")

            # Move to next hour
            current_hour += timedelta(hours=1)

    logger.info(f"Successfully processed partition. Inserted {total_rows_inserted} rows.")

    return dg.MaterializeResult(
        metadata={
            "run_id": dg.MetadataValue.text(run_id),
            "table_name": APP_AUDIO_ACTIONS_TABLE_NAME,
            "partition_time_window_start": dg.MetadataValue.text(partition_start.isoformat()),
            "partition_time_window_end": dg.MetadataValue.text(partition_end.isoformat()),
            "dagster/row_count": total_rows_inserted if not is_multi_partition_range else 0,
        },
    )