import argparse
import json
import stripe
import logging
from dataclasses import dataclass, field, asdict
import asyncio
import threading
import os
from dotenv import load_dotenv
from datetime import datetime
from snowflake.snowpark import Session


logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)

logging.getLogger("stripe").setLevel(logging.WARNING)

SEARCHED_EMAILS = set()
EMAIL_LOCK = threading.Lock()


load_dotenv()
STRIPE_API_KEY = os.getenv("STRIPE_API_KEY")
stripe.api_key = STRIPE_API_KEY


# Connection parameters from environment variables
SNOWFLAKE_CONNECTION_PARAMETERS = {
    "account": "fu90569.us-east-2.aws",
    "user": os.getenv("SNOWFLAKE_USERNAME"),
    "password": os.getenv("SNOWFLAKE_PASSWORD"),
    "role": "ACCOUNTADMIN",
    "warehouse": "SUNO_PROD_X_SMAL",
    "database": "SUNO_PROD",
    "schema": "PROD",  # Default to PUBLIC if not specified
}


@dataclass
class BillingAddress:
    city: str
    country: str
    line1: str
    line2: str
    postal_code: str
    state: str


@dataclass
class StripeCustomer:
    id: str
    email: str
    billing_address: BillingAddress | None = None
    card_fingerprints: list[str] = field(default_factory=list)
    related_customer_ids: list[str] = field(default_factory=list)

    def __hash__(self):
        return hash(self.id)  # Only use the unique identifier for hashing

    def __eq__(self, other):
        if not isinstance(other, StripeCustomer):
            return False
        return self.id == other.id  # Two customers are equal if they have the same ID


async def find_related_customers_by_payment_details(
    seed_customer_id,
) -> tuple[set[str], set[str]]:
    # returns a tuple of (related_customer_ids, payment_fingerprints)

    related_customers = set()

    # Use asyncio.to_thread to run blocking Stripe operations in a separate thread
    charges = await asyncio.to_thread(stripe.Charge.list, customer=seed_customer_id)

    payment_fingerprints = set()
    for charge in charges.data:
        if hasattr(charge, "payment_method_details"):
            source_type = charge.payment_method_details.type

            if source_type != "card":  # only handle card payments for now
                continue

            payment_details = getattr(charge.payment_method_details, source_type, None)

            if payment_details:
                fingerprint = getattr(payment_details, "fingerprint", None)
                payment_fingerprints.add(fingerprint)

    query_components = []

    for fp in payment_fingerprints:
        _query = f"payment_method_details.card.fingerprint:'{fp}'"
        query_components.append(_query)

    if not query_components:
        return (related_customers, payment_fingerprints)

    full_query = " OR ".join(query_components)

    try:
        similar_charges = await asyncio.to_thread(
            stripe.Charge.search, query=full_query
        )
        customers_for_fingerprint = {charge.customer for charge in similar_charges.data}

        related_customers.update(
            cust for cust in customers_for_fingerprint if cust != seed_customer_id
        )

    except stripe.error.InvalidRequestError as e:
        logging.error(f"Search query failed: {e}")

    return (related_customers, payment_fingerprints)


# async def find_related_customers_by_email(seed_customer_id):
async def find_related_customers_by_email(customer: stripe.Customer):
    related_customers = set()

    # Get the seed customer's email
    # customer = await asyncio.to_thread(stripe.Customer.retrieve, seed_customer_id)

    if (
        not hasattr(customer, "email")
        or not customer.get("email")
        or customer.get("email") == ""
    ):
        logging.info(f"No email found for customer {customer.id}")
        return related_customers

    email = customer["email"]

    with EMAIL_LOCK:
        if email in SEARCHED_EMAILS:
            logging.info(f"Already searched email {email}")
            return related_customers

    patterns = [
        f'email:"{email}"',  # Exact match
        # f'email:"{email_base}*@{domain}"',  # Variations with same base and domain
        # f'email:"*{email_base}*@{domain}"',  # More variations including prefixes
    ]

    # Search for each pattern
    for pattern in patterns:
        try:
            results = await asyncio.to_thread(stripe.Customer.search, query=pattern)
            for found_customer in results.auto_paging_iter():
                if found_customer.id != customer.id:
                    related_customers.add(found_customer.id)
                    logging.info(
                        f"Found related customer through email pattern '{pattern}': "
                        f"{found_customer.id} ({found_customer.email})"
                    )
        except stripe.error.InvalidRequestError as e:
            logging.error(f"Search query failed for pattern '{pattern}': {e}")
            continue

    with EMAIL_LOCK:
        SEARCHED_EMAILS.add(email)

    return related_customers


async def get_customer_info_and_find_related(customer_id) -> StripeCustomer:
    customer = await asyncio.to_thread(
        stripe.Customer.retrieve, customer_id, expand=["sources"]
    )

    # run payment and email searches in parallel
    (
        (related_customer_ids_by_payments, payment_fingerprints),
        related_by_email,
    ) = await asyncio.gather(
        find_related_customers_by_payment_details(customer_id),
        find_related_customers_by_email(customer),
    )

    # get billing address
    billing_address = customer.get("address")
    if billing_address:
        billing_address = BillingAddress(
            city=billing_address.get("city"),
            country=billing_address.get("country"),
            line1=billing_address.get("line1"),
            line2=billing_address.get("line2"),
            postal_code=billing_address.get("postal_code"),
            state=billing_address.get("state"),
        )
    #     logging.info(f"Billing address: {billing_address}")
    # else:
    #     logging.info("No billing address found")

    return StripeCustomer(
        id=customer.id,
        email=customer.get("email", ""),
        billing_address=billing_address,
        card_fingerprints=list(payment_fingerprints),
        related_customer_ids=list(related_customer_ids_by_payments | related_by_email),
    )


def load_file(file_path):
    with open(file_path, "r") as file:
        return [s.strip() for s in file.readlines()]


def create_session():
    try:
        # Create Snowflake session
        session = Session.builder.configs(SNOWFLAKE_CONNECTION_PARAMETERS).create()
        print("Successfully connected to Snowflake!")
        session.sql("USE WAREHOUSE SUNO_PROD_X_SMAL").collect()
        return session
    except Exception as e:
        print(f"Error connecting to Snowflake: {str(e)}")
        raise


def get_user_ids_by_stripe_customer_id(session, stripe_customer_ids):
    customer_id_strs = [f"'{customer_id}'" for customer_id in stripe_customer_ids]
    query = f"""
        SELECT user_id
        FROM suno_prod.prod.rds_discord_info
        WHERE stripe_customer_id IN ({", ".join(customer_id_strs)})
        AND stripe_customer_id IS NOT NULL
    """
    results = session.sql(query).collect()
    # print(results)
    return [row["USER_ID"] for row in results]


def fetch_suno_user_ids_from_snowflake(stripe_customer_ids):
    logging.info(
        f"Fetching user ids for {len(stripe_customer_ids)} stripe customer ids from snowflake"
    )
    session = create_session()
    user_ids = get_user_ids_by_stripe_customer_id(session, stripe_customer_ids)
    logging.info(f"Found {len(user_ids)} user ids in snowflake")
    return user_ids


async def async_main():
    parser = argparse.ArgumentParser(description="Find related Stripe customers")
    parser.add_argument(
        "--customer-ids",
        dest="customer_ids",
        required=False,
        nargs="+",
        help="One or more root customer IDs to start the search from",
    )
    parser.add_argument(
        "--customer-ids-file",
        dest="customer_ids_file",
        required=False,
        help="File containing customer IDs to process",
    )
    parser.add_argument(
        "--output-prefix",
        dest="output_prefix",
        required=True,
        help="Prefix for the output files",
    )

    args = parser.parse_args()

    if not args.customer_ids and not args.customer_ids_file:
        print("You must provide either --customer-ids or --customer-ids-file")
        return

    stripe_customer_ids = set()

    if args.customer_ids_file:
        stripe_customer_ids.update(load_file(args.customer_ids_file))

    if args.customer_ids:
        stripe_customer_ids.update(args.customer_ids)

    customers_to_process = set(stripe_customer_ids)
    processed_customers = set()
    all_related_customers = set(stripe_customer_ids)

    BATCH_SIZE = 10

    processed_customer_objects: set[StripeCustomer] = set()

    while customers_to_process:
        logging.info(f"Processing {len(customers_to_process)} customers")
        logging.info(f"Processed {len(processed_customers)} customers")
        logging.info(
            f"Found total related customers {len(all_related_customers)} customers"
        )

        batch = set()
        for _ in range(min(BATCH_SIZE, len(customers_to_process))):
            if customers_to_process:
                customer_id = customers_to_process.pop()
                if customer_id not in processed_customers:
                    batch.add(customer_id)

        if not batch:
            continue

        # Process customers in parallel
        results = await asyncio.gather(
            *(get_customer_info_and_find_related(customer) for customer in batch)
        )

        # Safely update sets with results
        for customer_id, customer in zip(batch, results):
            # Mark current customer as processed
            processed_customers.add(customer_id)
            processed_customer_objects.add(customer)

            # add to global related set
            all_related_customers.update(customer.related_customer_ids)
            all_related_customers.add(customer.id)  # probably redundant

            # Add newly found customers to processing queue
            new_to_process = {
                c
                for c in customer.related_customer_ids
                if c not in processed_customers and c not in customers_to_process
            }
            customers_to_process.update(new_to_process)

            logging.info(
                f"Found {len(new_to_process)} new related customers for {customer.id} {customer.email}"
            )

    logging.info(f"Total processed customers: {len(processed_customers)}")
    logging.info(f"Total related customers found: {len(all_related_customers)}")
    logging.info(f"All related customer_ids: {all_related_customers}")

    output_prefix = args.output_prefix

    current_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

    output_prefix_with_time = f"{output_prefix}_{current_time}_"

    customer_id_filename = f"{output_prefix_with_time}stripe_customer_ids.txt"
    email_filename = f"{output_prefix_with_time}stripe_customer_emails.txt"
    fingerprint_filename = f"{output_prefix_with_time}stripe_card_fingerprints.txt"

    stripe_customer_json_filename = (
        f"{output_prefix_with_time}stripe_customer_objects.jsonl"
    )

    user_id_filename = f"{output_prefix_with_time}suno_user_ids.txt"

    all_customer_fingerprints = set()
    distinct_emails = set()

    logging.info(f"Writing {len(processed_customer_objects)} customers to files")

    logging.info(
        f"Writing {len(processed_customer_objects)} customer ids to {customer_id_filename}"
    )
    with open(customer_id_filename, "w") as customer_ids_file:
        for customer in processed_customer_objects:
            all_customer_fingerprints.update(customer.card_fingerprints)

            if customer.email:
                distinct_emails.add(customer.email)

            customer_ids_file.write(customer.id + "\n")

    logging.info(f"Writing {len(distinct_emails)} distinct emails to {email_filename}")
    with open(email_filename, "w") as emails_file:
        for email in distinct_emails:
            emails_file.write(email + "\n")

    logging.info(
        f"Writing {len(all_customer_fingerprints)} card fingerprints to {fingerprint_filename}"
    )
    with open(fingerprint_filename, "w") as fingerprints_file:
        for fingerprint in all_customer_fingerprints:
            fingerprints_file.write(fingerprint + "\n")

    logging.info(
        f"Writing {len(processed_customer_objects)} stripe customers JSON objects to {stripe_customer_json_filename}"
    )
    with open(stripe_customer_json_filename, "w") as f:
        for customer in processed_customer_objects:
            f.write(json.dumps(asdict(customer)) + "\n")

    suno_user_ids = fetch_suno_user_ids_from_snowflake(all_related_customers)

    logging.info(f"Writing {len(suno_user_ids)} suno user ids to {user_id_filename}")
    with open(user_id_filename, "w") as f:
        for user_id in suno_user_ids:
            f.write(str(user_id) + "\n")

    logging.info("All done -- happy hunting!")


def main():
    asyncio.run(async_main())


if __name__ == "__main__":
    """
    See README.md for more detailed instructions.

    python stripe_find_related_accounts.py --output-prefix <prefix> --customer-ids <stripe_customer_id>

    Output prefix is used to name the output files, datetime is also added for clarity.

    You can also pass in a file of customer ids to process:
    python stripe_find_related_accounts.py --output-prefix <prefix> --customer-ids-file <file>
    """
    main()