import os
from typing import List

from snowflake.snowpark.session import Session


def get_create_suno_table_command(
    af_session: Session,
    table: str,
) -> str:
    # Get column definitions from APPSFLYER table
    cols_query = f"""
        SELECT column_name, data_type, character_maximum_length, numeric_precision, numeric_scale
        FROM APPSFLYER.INFORMATION_SCHEMA.COLUMNS
        WHERE table_schema = 'SNOWFLAKE'
        AND table_name = '{table}'
        ORDER BY ordinal_position
    """
    cols = af_session.sql(cols_query).collect()

    # Build column definitions string
    col_defs = []
    for col in cols:
        col_type = col["DATA_TYPE"]
        if col_type == "TEXT":
            if col["CHARACTER_MAXIMUM_LENGTH"]:
                col_type = f"VARCHAR({col['CHARACTER_MAXIMUM_LENGTH']})"
        elif col_type == "NUMBER":
            p = col["NUMERIC_PRECISION"]
            s = col["NUMERIC_SCALE"]
            if p is not None and s is not None:
                col_type = f"NUMBER({p},{s})"
        col_defs.append(f"{col['COLUMN_NAME']} {col_type}")

    # Construct SQL command to create the table
    dst_table_name = f"AF_{table.upper().replace('_VIEW', '')}"
    columns_str = ",\n    ".join(col_defs)
    create_table_sql = f"""
CREATE TRANSIENT TABLE SUNO_PROD.PROD.{dst_table_name} (
    {columns_str}
);
"""

    return create_table_sql


def get_alter_table_command(
    af_session: Session,
    table: str,
    cols_to_add: List[str],
) -> str:
    if not cols_to_add:
        return ""

    # Get column definitions from APPSFLYER table
    query = f"""
        SELECT column_name, data_type, character_maximum_length, numeric_precision, numeric_scale
        FROM APPSFLYER.INFORMATION_SCHEMA.COLUMNS
        WHERE table_schema = 'SNOWFLAKE'
        AND table_name = '{table}'
        AND column_name IN ({",".join([f"'{col}'" for col in cols_to_add])})
        ORDER BY ordinal_position
    """
    cols_info = af_session.sql(query).collect()

    # Generate ALTER TABLE command
    dst_table_name = f"AF_{table.upper().replace('_VIEW', '')}"
    alter_cmd = f"""
ALTER TABLE SUNO_PROD.PROD.{dst_table_name} ADD COLUMN
"""
    for col in cols_info:
        col_name = col["COLUMN_NAME"]
        data_type = col["DATA_TYPE"]

        # Handle different data types
        if data_type == "VARCHAR":
            max_length = col["CHARACTER_MAXIMUM_LENGTH"]
            col_def = f"VARCHAR({max_length})"
        elif data_type in ["NUMBER", "DECIMAL", "NUMERIC"]:
            precision = col["NUMERIC_PRECISION"]
            scale = col["NUMERIC_SCALE"]
            col_def = f"NUMBER({precision},{scale})"
        else:
            col_def = data_type

        alter_cmd += f"    {col_name} {col_def},\n"
    alter_cmd = alter_cmd.rstrip(",\n") + ";\n"

    return alter_cmd


def create_task_file(session: Session, af_table_name: str):
    """Create a task.sql file for a new AppsFlyer table
    Overwrites the file if it already exists.

    Args:
        session: Snowflake session
        af_table_name: Name of the table in APPSFLYER.SNOWFLAKE (e.g. CLICKS_VIEW)
    """
    # Convert CLICKS_VIEW to AF_CLICKS
    suno_table_name = f"AF_{af_table_name.replace('_VIEW', '')}"

    # Create the directory if it doesn't exist
    task_dir = f"tables/appsflyer/{suno_table_name.lower()}"
    os.makedirs(task_dir, exist_ok=True)

    task_content = f"""create or replace task SUNO_PROD.PROD.{suno_table_name}_DAILY_INSERT
    warehouse=SUNO_PROD_RDS_DAILY_X_SMALL
    schedule='USING CRON 0 11 * * * UTC'
as
    declare
        start_time TIMESTAMP;
        max_ingestion_time timestamp;
        row_count number;
    begin
        start_time := CURRENT_TIMESTAMP();
        SELECT COALESCE(MAX(_INGESTION_TIME), '0001-01-01')
            INTO :max_ingestion_time
        FROM SUNO_PROD.PROD.{suno_table_name};

        INSERT INTO SUNO_PROD.PROD.{suno_table_name} (
            {{columns}}
        )
        SELECT
            {{columns}}
        FROM APPSFLYER.SNOWFLAKE.{af_table_name}
        WHERE _INGESTION_TIME > :max_ingestion_time;

        SELECT COUNT(*)
            INTO :row_count
        FROM SUNO_PROD.PROD.{suno_table_name}
        WHERE _INGESTION_TIME > :max_ingestion_time;
        call TASK_MONITOR_INSERT_PROC('{suno_table_name}_DAILY_INSERT', :start_time, :row_count, 'SUCCESS', 'DAILY');
    end
;
"""

    # Get the columns from the APPSFLYER table
    columns_query = f"""
    SELECT COLUMN_NAME 
    FROM APPSFLYER.INFORMATION_SCHEMA.COLUMNS 
    WHERE TABLE_SCHEMA = 'SNOWFLAKE' 
    AND TABLE_NAME = '{af_table_name}'
    ORDER BY ORDINAL_POSITION;
    """

    columns = [row["COLUMN_NAME"] for row in session.sql(columns_query).collect()]
    columns_str = ",\n            ".join(columns)

    # Replace the {columns} placeholder with actual columns
    task_content = task_content.replace("{columns}", columns_str)

    # Write the task file
    task_file_path = f"{task_dir}/task.sql"
    with open(task_file_path, "w") as f:
        f.write(task_content)