import argparse
import os
import sys
from typing import Dict, List

sys.path.append(".")

from snowflake.snowpark.session import Session

from utils.appsflyer.table_commands import (
    create_task_file,
    get_alter_table_command,
    get_create_suno_table_command,
)
from utils.common import get_snowflake_session, get_table_columns, get_table_names

SNOWFLAKE_WAREHOUSE = "SUNO_PROD_X_SMAL"


def parse_args():
    """Parse command line arguments"""
    parser = argparse.ArgumentParser(
        description="Compare columns between APPSFLYER and SUNO_PROD tables"
    )
    parser.add_argument(
        "--apply-changes",
        action="store_true",
        help="Apply the changes to the database and repo",
    )
    return parser.parse_args()


# Procedure
# 1. get all the tables and cols from SUNO_PROD.PROD that start with AF_
# 2. get all the tables and cols from APPSFLYER.SNOWFLAKE
# 3. diff the table populations
# 4. if there are any table differences, generate a CREATE TRANSIENT TABLE command and print it out
# 5. for each table common to both databases, compare the column names
# 6. if there are differences in columns for a table, print them out
# 7. if APPSFLYER has extra columns, print the ALTER TABLE command to add them
# 8. if SUNO_PROD has extra columns, TBD. Don't expect this to happen.


def snowflake_proc_entry_point(session: Session):
    return main(
        session,
        apply_changes=False,
    )


def prompt_and_create_task_files_and_resume(session: Session, af_table_names: List[str]):
    response = input(
        "Do you want to (re)create the Snowflake task files for these tables in the repo? (y/n): "
    )
    if response.lower() in ["y", "yes"]:
        for af_table_name in af_table_names:
            print(f"(Re)creating {af_table_name} task file in the repo...")
            create_task_file(session, af_table_name)

    response = input("Do you want to start these tasks in Snowflake? (y/n): ")
    if response.lower() in ["y", "yes"]:
        for af_table_name in af_table_names:
            print(f"(Re)starting {af_table_name} task in Snowflake...")
            suno_table_name = f"AF_{af_table_name.replace('_VIEW', '')}"
            session.sql(f"ALTER TASK SUNO_PROD.PROD.{suno_table_name}_DAILY_INSERT RESUME").collect()
        print(f"Resumed {len(af_table_names)} task(s) in Snowflake")


def main(
    session: Session,
    apply_changes: bool = False,
):
    """
    Compare the tables and columns between APPSFLYER and SUNO_PROD.
    Return true if there are any differences. False otherwise

    Args:
        session: Snowflake session
        apply_changes: Run the commands to the database and update task files in the repo
    """
    # Get the available tables in each database
    print("Getting APPSFLYER and SUNO_PROD tables...")
    suno_table_names = get_table_names(session, "SUNO_PROD", "PROD", "AF_")
    af_table_names = get_table_names(session, "APPSFLYER", "SNOWFLAKE")

    suno_table_names_converted = [
        f"{table_name.replace('AF_', '')}_VIEW" for table_name in suno_table_names
    ]
    # Get the tables that are only in APPSFLYER but not in SUNO_PROD
    af_new_table_names = sorted(set(af_table_names) - set(suno_table_names_converted))
    if len(af_new_table_names) > 0:
        print(
            f"Found {len(af_new_table_names)} missing tables in SUNO_PROD database. "
            "Dumping CREATE TABLE commands for new tables to stdout"
        )
    else:
        print("No new tables in APPSFLYER database :)")

    # For each table that is only in APPSFLYER, get the create table command and print it out
    for table_name in af_new_table_names:
        create_table_command = get_create_suno_table_command(session, table_name)
        print(create_table_command)

    # Get user confirmation before creating SUNO_PROD tables in Snowflake
    if af_new_table_names and apply_changes:
        response = input(
            f"Do the commands for creating {', '.join(af_new_table_names)}"
            " in Snowflake look good? (y/n): "
        )
        if response.lower() in ["y", "yes"]:
            for table_name in af_new_table_names:
                print(f"Creating {table_name} table in Snowflake...")
                session.sql(create_table_command).collect()

        prompt_and_create_task_files_and_resume(session, af_new_table_names)

    # Print the tables that AppsFlyer dropped from their own APPSFLYER database
    af_missing_table_names = set(suno_table_names_converted) - set(af_table_names)
    if len(af_missing_table_names) > 0:
        print(
            "Missing tables in APPSFLYER database"
            f" (not sure why this would happen): {list(af_missing_table_names)}"
        )
    else:
        print("No missing tables in APPSFLYER database :)")

    print("Comparing columns between existing APPSFLYER and SUNO_PROD tables...")
    # For the tables that are in both databases, get the columns and compare them
    af_mutual_tables = sorted(set(suno_table_names_converted) & set(af_table_names))
    suno_mutual_tables = [f"AF_{table_name.replace('_VIEW', '')}" for table_name in af_mutual_tables]
    suno_table_cols = get_table_columns(session, "SUNO_PROD", "PROD", suno_mutual_tables)
    af_table_cols = get_table_columns(session, "APPSFLYER", "SNOWFLAKE", af_mutual_tables)

    # Standardize the column names to use the AF format
    suno_table_cols_af_format = {
        f"{table_name.replace('AF_', '')}_VIEW": val for table_name, val in suno_table_cols.items()
    }
    differences = compare_table_columns(suno_table_cols_af_format, af_table_cols)
    if len(differences) > 0:
        print("Dumping ALTER TABLE commands for tables with missing columns to stdout")

    else:
        print("No differences in columns between APPSFLYER and SUNO_PROD databases :)")

    # This contains table names in AF format
    af_table_names_with_diff_cols = sorted(differences.keys())
    add_col_to_suno_table_commands = {}
    for af_table_name in af_table_names_with_diff_cols:
        diff = differences[af_table_name]
        af_only_cols = diff["af_only_cols"]
        if af_only_cols:
            alter_table_command = get_alter_table_command(session, af_table_name, af_only_cols)
            add_col_to_suno_table_commands[af_table_name] = alter_table_command
            print(alter_table_command)

    # Get user confirmation before add columns to SUNO_PROD tables in Snowflake
    if add_col_to_suno_table_commands and apply_changes:
        suno_diff_cols_table_names = [
            f"AF_{table_name.replace('_VIEW', '')}" for table_name in af_table_names_with_diff_cols
        ]
        response = input(
            f"Run SQL commands in Snowflake to alter {', '.join(suno_diff_cols_table_names)}? (y/n): "
        )
        if response.lower() in ["y", "yes"]:
            for af_table_name in af_table_names_with_diff_cols:
                print(f"Altering {af_table_name} table in Snowflake...")
                session.sql(add_col_to_suno_table_commands[af_table_name]).collect()

        prompt_and_create_task_files_and_resume(session, af_table_names_with_diff_cols)

    # Print out any columns that AppsFlyer dropped in their own APPSFLYER database
    for af_table_name in af_table_names_with_diff_cols:
        diff = differences[af_table_name]
        suno_only_cols = diff["suno_only_cols"]
        if suno_only_cols:
            print(
                "Missing columns in APPSFLYER database"
                f" (not sure why this would happen): {list(suno_only_cols)}"
            )

    return len(af_new_table_names) > 0 or len(differences) > 0


def compare_table_columns(
    suno_tables: Dict[str, List[str]], af_tables: Dict[str, List[str]]
) -> Dict[str, Dict[str, List[str]]]:
    """Compare columns between two dictionaries of tables mapped to their columns

    Returns a dictionary {SUNO.PROD table name -> {suno_only_cols: [], af_only_cols: []}}
    """

    assert set(suno_tables.keys()) == set(af_tables.keys()), "Table set must be the same"

    # Compare columns for each common table
    differences = {}
    for table in suno_tables:
        suno_cols = suno_tables[table]
        af_cols = af_tables[table]

        # Find differences in both directions
        suno_only_cols = set(suno_cols) - set(af_cols)
        af_only_cols = set(af_cols) - set(suno_cols)

        if suno_only_cols or af_only_cols:
            differences[table] = {
                "suno_only_cols": suno_only_cols,
                "af_only_cols": af_only_cols,
            }

    return differences


if __name__ == "__main__":
    from dotenv import load_dotenv

    load_dotenv()

    SNOWFLAKE_CONFIGS = {
        "account": os.getenv("SNOWFLAKE_ACCOUNT"),
        "user": os.getenv("SNOWFLAKE_ACCOUNT_USER"),
        "password": os.getenv("SNOWFLAKE_ACCOUNT_PASSWORD"),
        "role": os.getenv("SNOWFLAKE_ACCOUNT_ROLE"),
    }

    args = parse_args()

    session = get_snowflake_session(
        snowflake_configs=SNOWFLAKE_CONFIGS,
        warehouse=SNOWFLAKE_WAREHOUSE,
        database="SUNO_PROD",
        schema="PROD",
    )
    main(session, args.apply_changes)