import argparse
import csv
import os
from typing import List

import pandas as pd
import snowflake.connector
import splitio
from dotenv import load_dotenv

IDENTITY_ID_ATTRIBUTE = "identity_id"

FEATURE_FLAGS = [
    'orchard_suite_show_audience_app',
    'audience_apply_analyst_role',
    'show_sme_data',
    'audience_allow_audience_export_to_file'
]

# Load environment variables from .env file
load_dotenv()


def fetch_ids_from_snowflake(emails: List[str]) -> List[List]:
    """Fetch IDs from Snowflake for provided emails.

    Args:
        emails (List[str]): List of email addresses.

    Returns:
        List[List]: List of lists containing IDs and corresponding emails.
    """
    snowflake_config = {
        'user': os.getenv('SNOWFLAKE_USER'),
        'password': os.getenv('SNOWFLAKE_PASSWORD'),
        'account': os.getenv('SNOWFLAKE_ACCOUNT'),
        'role': os.getenv('SNOWFLAKE_ROLE'),
        'warehouse': os.getenv('SNOWFLAKE_WAREHOUSE'),
        'database': os.getenv('SNOWFLAKE_DATABASE'),
        'schema': os.getenv('SNOWFLAKE_SCHEMA'),
        'passcode': os.getenv('SNOWFLAKE_PASSCODE'),
        'table': os.getenv('SNOWFLAKE_TABLE')
    }

    try:
        conn = snowflake.connector.connect(**snowflake_config)
        cursor = conn.cursor()

        query = f"""
        SELECT DISTINCT id, email
        FROM {snowflake_config['table']}
        WHERE email IN ({','.join(["'" + email + "'" for email in emails])})
        """

        cursor.execute(query)
        results = [[row[0], row[1]] for row in cursor.fetchall()]

        cursor.close()
        conn.close()

        return results

    except snowflake.connector.errors.ProgrammingError as e:
        print("Snowflake error:", e)
        return []
    except Exception as e:
        print("Error:", e)
        return []


def check_treatments(ids: List[List], feature_flags: List[str], attribute: str) -> List[List[int]]:
    """Check treatments using IDs retrieved from Snowflake.

    Args:
        ids (List[List]): List of lists containing IDs and corresponding emails.
        feature_flags (List[str]): List of feature flags.
        attribute (str): Attribute to use for Split.io client.

    Returns:
        List[List[int]]: List of lists containing treatments for each ID.
    """
    factory = splitio.get_factory(os.getenv('SPLIT_API_KEY'))
    factory.block_until_ready(5)
    splitio_client = factory.client()

    treatments = []

    for id in ids:
        user_treatments = []
        for feature_flag in feature_flags:
            treatment = splitio_client.get_treatment(
                key=attribute,
                feature_flag=feature_flag,
                attributes={attribute: id[0]}
            )
            user_treatments.append(1 if treatment == "on" else 0)

        treatments.append(user_treatments)

    return treatments


def parse_emails_from_csv(csv_file: str) -> List[str]:
    """Parse CSV file and extract emails.

    Args:
        csv_file (str): Path to CSV file containing email addresses without header.

    Returns:
        List[str]: List of email addresses.
    """
    emails = []
    with open(csv_file, 'r') as file:
        reader = csv.reader(file)
        for row in reader:
            if row:
                emails.extend(row)
    return emails


def create_dataframe(ids: List[List], treatments: List[List[int]], column_names: List[str]) -> pd.DataFrame:
    """Create Pandas DataFrame from results and treatments.

    Args:
        ids (List[List]): List of lists containing IDs and corresponding emails.
        treatments (List[List[int]]): List of lists containing treatments for each ID.
        column_names (List[str]): List of column names for the DataFrame.

    Returns:
        pd.DataFrame: Pandas DataFrame.
    """
    rows = []
    for id, treatment in zip(ids, treatments):
        row = id + treatment
        rows.append(row)

    df = pd.DataFrame(rows, columns=column_names)
    return df


def save_to_csv(df: pd.DataFrame, filename: str):
    """Save DataFrame to CSV file.

    Args:
        df (pd.DataFrame): Pandas DataFrame.
        filename (str): Filename to save the DataFrame.
    """
    df.to_csv(filename, index=False)


def main():
    parser = argparse.ArgumentParser(description="Fetch IDs and check treatments")
    parser.add_argument("-e", "--emails", metavar="EMAIL1,EMAIL2,...", type=str,
                        help="Comma-separated list of email addresses")
    parser.add_argument("-f", "--csv_file", metavar="FILE", type=str,
                        help="Path to CSV file containing email addresses without header")
    parser.add_argument("-s", "--save", metavar="FILENAME", type=str, nargs='?',
                        const="results.csv", default=None,
                        help="Save results to CSV file with the specified filename")

    args = parser.parse_args()

    if args.emails:
        emails = args.emails.split(',')
    elif args.csv_file:
        emails = parse_emails_from_csv(args.csv_file)
    else:
        print("Error: Please provide either comma-separated emails (-e) or a CSV file (-f)")
        return

    ids = fetch_ids_from_snowflake(emails)

    if not ids:
        print("No matching IDs found for the provided emails.")
        return

    treatments = check_treatments(ids, FEATURE_FLAGS, IDENTITY_ID_ATTRIBUTE)
    colnames = ['ID', 'EMAIL'] + FEATURE_FLAGS
    df = create_dataframe(ids, treatments, colnames)

    if args.save is not None:
        save_to_csv(df, args.save)
        print(f"Results saved to '{args.save}'")
    else:
        print(df)


if __name__ == "__main__":
    main()