import argparse
import stripe
from collections import defaultdict, deque
import re
import logging
from dataclasses import dataclass, field, asdict

import os
from dotenv import load_dotenv
import json
import asyncio


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

load_dotenv(dotenv_path=".env")

STRIPE_API_KEY = os.getenv("STRIPE_API_KEY")

logging.info(f"Using Stripe API key: {STRIPE_API_KEY}")

stripe.api_key = STRIPE_API_KEY


def read_customer_ids(filepath):
    """Read customer IDs from a file, one ID per line."""
    with open(filepath, "r") as f:
        return [line.strip() for line in f if line.strip()]


@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)


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

    email = customer.get("email")
    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"),
        )

    # Get charges for the seed customer
    charges = await asyncio.to_thread(stripe.Charge.list, customer=customer_id)

    # Extract payment method details from each charge
    payment_fingerprints = set()
    for charge in charges.data:
        if hasattr(charge, "payment_method_details"):
            source_type = charge.payment_method_details.type

            # Skip all besides credit card today
            if source_type != "card":
                continue

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

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

    return StripeCustomer(
        id=customer.id,
        email=email,
        billing_address=billing_address,
        card_fingerprints=list(payment_fingerprints),
    )


async def process_customer_batch(batch: list[str]) -> list[StripeCustomer]:
    tasks = [get_stripe_customer_info(customer_id) for customer_id in batch]
    return await asyncio.gather(*tasks)


def process_all_customer_batches(customer_ids: list[str], batch_size: int = 10):
    customers = []
    total_fingerprints = 0

    for i in range(0, len(customer_ids), batch_size):
        batch = customer_ids[i : i + batch_size]
        batch_results = asyncio.run(process_customer_batch(batch))
        customers.extend(batch_results)

        for customer in batch_results:
            total_fingerprints += len(customer.card_fingerprints)

        logging.info(f"Processed batch of {len(batch)} customers")
        logging.info(f"Processed {len(customers)} / {len(customer_ids)} customers")
        logging.info(f"Found {total_fingerprints} card fingerprints so far")

    return customers


def op_get_customer_info_and_fingerprints(args):
    customer_ids = read_customer_ids(args.customer_ids_file)
    if args.test:
        customer_ids = customer_ids[:100]
    logging.info(f"Processing {len(customer_ids)} customer IDs")

    BATCH_SIZE = 35

    customers = process_all_customer_batches(customer_ids, BATCH_SIZE)

    # Write customers to JSON file, one per line
    with open(f"{args.output_file_prefix}_customers.jsonl", "w") as f:
        for customer in customers:
            json.dump(asdict(customer), f)
            f.write("\n")

    logging.info(f"Wrote {len(customers)} customers to {args.output_file_prefix}")

    fingerprints = set()
    for customer in customers:
        if customer.card_fingerprints:
            fingerprints.update(customer.card_fingerprints)

    fingerprint_file = f"{args.output_file_prefix}_fingerprints.txt"
    with open(fingerprint_file, "w") as f:
        for fingerprint in fingerprints:
            f.write(f"{fingerprint}\n")

    logging.info(f"Wrote {len(fingerprints)} fingerprints to {fingerprint_file}")


def op_sample_customer_urls(args):
    customer_ids = read_customer_ids(args.customer_ids_file)

    import random

    random.shuffle(customer_ids)
    sample = customer_ids[:50]

    customers = process_all_customer_batches(sample, 10)

    logging.info("=======================================")
    logging.info(f"Sampled {len(sample)} random customer IDs")
    logging.info("=======================================")
    for customer in customers:
        logging.info(f"Customer: {customer}")
        logging.info(
            f"Customer Stripe URL: https://dashboard.stripe.com/customers/{customer.id}"
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process Stripe customer information")
    parser.add_argument("--op", required=True, help="Operation to perform")
    parser.add_argument(
        "--customer-ids-file",
        required=True,
        help="Path to file containing Stripe customer IDs (one per line)",
    )
    parser.add_argument(
        "--output-file-prefix",
        required=False,
        help="Path to output file",
    )
    parser.add_argument(
        "--test",
        action="store_true",
        required=False,
        default=False,
        help="Only process the first 100 customer IDs",
    )

    args = parser.parse_args()

    if args.op == "get-customers-and-fingerprints":
        if not args.output_file_prefix:
            raise ValueError("Output file prefix is required for this operation")
        op_get_customer_info_and_fingerprints(args)
    elif args.op == "sample-customer-urls":
        op_sample_customer_urls(args)
    else:
        raise ValueError(f"Unknown operation: {args.op}")