import argparse
import datetime
import stripe
from stripe.error import InvalidRequestError, APIError
import logging
from dataclasses import dataclass, field
import sys

import os
from dotenv import load_dotenv
import json


from urllib.parse import quote_plus
from django.forms.models import model_to_dict
from django.core.serializers.json import DjangoJSONEncoder

from markdown_pdf import MarkdownPdf, Section

from pypdf import PdfWriter

import concurrent.futures
from typing import Tuple


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

load_dotenv(dotenv_path=".env")

STRIPE_API_KEY = os.getenv("STRIPE_WRITE_API_KEY")

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

stripe.api_key = STRIPE_API_KEY


project_root = os.path.abspath(
    os.path.join(os.path.dirname(__file__), "../../studio_api")
)
os.chdir(project_root)
sys.path.append(project_root)

os.environ.setdefault("DJANGO_SETTINGS_MODULE", "studio_api.settings")
os.environ["DJANGO_ALLOW_ASYNC_UNSAFE"] = (
    "true"  # https://docs.djangoproject.com/en/4.1/topics/async/#async-safety
)

database_url = f"postgresql://{os.getenv('DATABASE_USER')}:{quote_plus(os.getenv('DATABASE_PASSWORD'))}@{os.getenv('DATABASE_HOST')}/{os.getenv('DATABASE_NAME')}"
replica_database_url = database_url

print("Using database:")
print(f"- host: {os.getenv('DATABASE_HOST')}")
print(f"- username: {os.getenv('DATABASE_USER')}")
print(
    f"- password: {'*' * len(os.getenv('DATABASE_PASSWORD')) if os.getenv('DATABASE_PASSWORD') else 'None'}"
)
print(f"- database: {os.getenv('DATABASE_NAME')}")

if not database_url or not replica_database_url:
    raise ValueError(
        "DATABASE_URL and REPLICA_DATABASE_URL must be set in environment variables"
    )

os.environ["DATABASE_URL"] = database_url
os.environ["REPLICA_DATABASE_URL"] = replica_database_url

import django

django.setup()

from studio_api.bots.models import DiscordInfo, GeneratedClip, UserStats


DISPUTE_TEMPLATE_MARKDOWN = """
We believe this dispute is illegitimate because we collected legitimate payment information from the customer linked to the email address and IP address they used to sign in to our platform, and we provided the customer with the subscription they paid for as well as digital goods via that subscription that they created and have irrevokable access to via downloads.

Our platform Suno is used by customers to create songs.
**Customer "{customer_name}" ({email}) purchased a subscription from us on {subscription_date} for ${subscription_amount}.** 
This is a {subscription_type} giving the ability to create hundreds of songs on our platform.
A typical subscription is $10/month for 500+ song creations or $30/month for 2000+ song creations. 

The customer received the subscription they paid for relevant to the TOS they agree to.
The customer paid for monthly credits to create songs on our platform. 

We validated the customer's credit card and ensured a matching email address to the one they authorized to our service with to ensure the validity of the payment and link the customer's identity between their payment and account. We require that the email address on the Stripe transaction matches an authenticated and verified email address that the customer signed up with via sign in via an OAuth provider like Google or Microsoft.

The customer's credit card was checked with the following results, the best available to us for the given card:
- CVC: {cvc_result}
- Address Line 1: {address_line1_result}
- Address Postal Code: {address_postal_code_result}

We can see they connected from an IP address that matches the stripe payment IP of: {ip_address}

In our terms of service we state that payments are non-refundable:

> "Refunds and Cancellations: Payments made by you hereunder are final and non-refundable, unless otherwise determined by Suno."

"""

SONG_COUNT_TEMPLATE = """
**The customer generated {song_count} songs (2 versions per song) in the month after the charge, which they have irrevokable access to now via downloads. 
They have generated approximately {total_songs} songs in total representing meaningful activity on our platform. 
These song creations cost us money to create by consuming compute resources we pay for.**
Note: This the two are counted slightly differently, so may not match exactly.
"""

DELETED_ACCOUNT_TEMPLATE = """
The customer deleted their Suno account, this is likely a tactic by a fraudster to make it harder for us to analyze their account after the fact.
As a result we'll have limited information about their account and activity, but given the correlation between email address, IP address, and a valid payment method we believe the dispute is illegitimate.
"""


def get_dispute(dispute_id):
    """
    Fetch a specific dispute by its ID.

    Args:
        dispute_id (str): The ID of the dispute to retrieve

    Returns:
        dict: The dispute object if found

    Raises:
        InvalidRequestError: If the dispute doesn't exist or other request issues
    """
    try:
        dispute = stripe.Dispute.retrieve(dispute_id)
        logging.info(f"Successfully retrieved dispute: {dispute_id}")
        return dispute
    except InvalidRequestError as e:
        logging.error(f"Error retrieving dispute {dispute_id}: {str(e)}")
        raise
    except APIError as e:
        logging.error(f"Stripe API error: {str(e)}")
        raise


@dataclass
class DisputeInformation:
    dispute: stripe.Dispute
    customer: stripe.Customer | None = None
    charge: stripe.Charge | None = None
    clips: list[GeneratedClip] = field(default_factory=list)
    cvc_result: str | None = None
    address_line1_result: str | None = None
    address_postal_code_result: str | None = None
    user_stats: UserStats | None = None
    account_was_deleted: bool = False


def get_song_list_markdown(clips: list[GeneratedClip]) -> str:
    base = f"""
Below is a list of songs created by the customer on Suno.com in the relevant subscription period.

### {len(clips)}+ Songs created by the customer within the month after the charge:

Note: Each song listed represents 2 versions of a song provided by Suno to the customer, list of songs may be truncated due to duplicates and PDF size limits.

"""
    original_len = len(clips)
    clips = clips[:100]
    truncated = len(clips) != original_len
    last_clip_string = ""

    for clip in clips:
        title = clip.title or "[No title]"
        clip_time_string = f"- {title} ({clip.created_at.strftime('%m/%d/%y %H:%M')})\n"

        if clip_time_string != last_clip_string:
            base += clip_time_string
            last_clip_string = clip_time_string

    delta = original_len - len(clips)
    if truncated:
        base += f"- {delta} additional songs not listed due to PDF size limits\n"

    return base


def get_user_stats_markdown(user_stats: UserStats) -> str:
    return f"""
We can see the user has activity on our platform as represented by their overall usage statistics:

```
{to_json(user_stats)}
```
"""


def to_json(obj):
    return json.dumps(model_to_dict(obj), indent=2, cls=DjangoJSONEncoder)


def get_dispute_markdown(dispute_info: DisputeInformation) -> str:
    dispute_dollars = int(dispute_info.charge.amount / 100)

    match dispute_dollars:
        case 10:
            subscription_type = "Suno.com monthly pro subscription"
        case 30:
            subscription_type = "Suno.com monthly premier subscription"
        case 4, 8, 16:
            subscription_type = "Suno.com credit top-up (more song creations)"
        case 96:
            subscription_type = "Suno.com yearly pro subscription"
        case 288:
            subscription_type = "Suno.com yearly premier subscription"
        case _:
            subscription_type = "subscription"

    result = DISPUTE_TEMPLATE_MARKDOWN.format(
        customer_name=dispute_info.customer.name,
        email=dispute_info.customer.email,
        subscription_date=datetime.datetime.fromtimestamp(dispute_info.charge.created),
        subscription_amount=dispute_info.charge.amount / 100,
        # song_count=len(dispute_info.clips),
        cvc_result=dispute_info.cvc_result,
        address_line1_result=dispute_info.address_line1_result,
        address_postal_code_result=dispute_info.address_postal_code_result,
        subscription_type=subscription_type,
        ip_address=dispute_info.dispute.evidence.customer_purchase_ip,
        # user_stats=to_json(dispute_info.user_stats),
    )

    if dispute_info.account_was_deleted:
        result += DELETED_ACCOUNT_TEMPLATE
    else:
        result += SONG_COUNT_TEMPLATE.format(
            song_count=len(dispute_info.clips),
            total_songs=dispute_info.user_stats.total_clips
            if dispute_info.user_stats
            else len(dispute_info.clips),
        )

        if dispute_info.user_stats:
            result += get_user_stats_markdown(dispute_info.user_stats)

        result += get_song_list_markdown(dispute_info.clips)

    return result


def sprint(obj):
    print(json.dumps(obj, indent=2))


def dprint(obj):
    print(json.dumps(model_to_dict(obj), indent=2, cls=DjangoJSONEncoder))


def hydrate_dispute_information(dispute: stripe.Dispute):
    # customer = dispute.customer
    # sprint(dispute)
    # print(customer)
    print("Hydrating dispute information")
    charge_id = dispute.charge
    charge = stripe.Charge.retrieve(charge_id)
    # sprint(charge)

    checks = charge.payment_method_details.card.checks

    cvc_result = checks.cvc_check
    address_line1_result = checks.address_line1_check
    address_postal_code_result = checks.address_postal_code_check

    customer_id = charge.customer
    customer = stripe.Customer.retrieve(customer_id)
    # sprint(customer)

    try:
        discord_info = DiscordInfo.objects.get(stripe_customer_id=customer.id)
        # dprint(discord_info)
    except DiscordInfo.DoesNotExist:
        discord_info = None

    charge_created_at_seconds = charge.created
    charge_created_at = datetime.datetime.fromtimestamp(charge_created_at_seconds)
    charge_month_end_approx = charge_created_at + datetime.timedelta(days=32)

    if discord_info:
        clips = GeneratedClip.objects.filter(
            user=discord_info.user,
            created_at__gte=charge_created_at,
            created_at__lte=charge_month_end_approx,
        )
    else:
        clips = []

    print(f"Found {len(clips)} clips in the month after the charge.")
    # for clip in clips:
    # print(f"Clip {clip.id} created at {clip.created_at} -- {clip.title}")

    user_stats = None  # Initialize user_stats with a default value
    try:
        if discord_info:
            user_stats = UserStats.objects.get(user=discord_info.user)
    except UserStats.DoesNotExist:
        print(f"User stats for user {discord_info.user.id} not found")
        pass

    return DisputeInformation(
        dispute=dispute,
        customer=customer,
        charge=charge,
        clips=clips,
        cvc_result=cvc_result,
        address_line1_result=address_line1_result,
        address_postal_code_result=address_postal_code_result,
        user_stats=user_stats,
        account_was_deleted=discord_info is None,
    )


def create_and_save_pdf(markdown_content: str, output_file: str):
    print(f"Creating PDF and saving to {output_file}")
    pdf = MarkdownPdf()
    pdf.add_section(Section(markdown_content, toc=False, paper_size="A4"))
    pdf.save(output_file)
    print(f"PDF saved to {output_file}")
    print("compressing PDF")
    writer = PdfWriter(clone_from=output_file)
    writer.compress_identical_objects(remove_identicals=True, remove_orphans=True)
    for page in writer.pages:
        page.compress_content_streams()
    with open(output_file, "wb") as f:
        writer.write(f)
    print(f"Compressed PDF saved to {output_file}")


def get_usage_markdown(dispute_info: DisputeInformation) -> str:
    base_info = f"""
email_address: {dispute_info.customer.email}
ip_address: {dispute_info.dispute.evidence.customer_purchase_ip}

"""
    if dispute_info.user_stats:
        return base_info + get_user_stats_markdown(dispute_info.user_stats)
    else:
        return base_info


def update_dispute_details(
    dispute_id: str, dispute_info: DisputeInformation, file_path: str, submit: bool
):
    """
    Upload a file to Stripe and update the dispute evidence with that file.

    Args:
        dispute_id (str): The ID of the dispute to update
        file_path (str): Path to the file to upload as evidence
    """
    logging.info(f"Updating dispute {dispute_id} with file {file_path}")

    # First, get the dispute to ensure it exists
    # dispute = get_dispute(dispute_id)

    usage_markdown = get_usage_markdown(dispute_info)

    # Upload the file to Stripe
    try:
        with open(file_path, "rb") as file:
            file_upload = stripe.File.create(
                purpose="dispute_evidence", file=file, file_link_data={"create": True}
            )
        logging.info(f"File uploaded successfully with ID: {file_upload.id}")

        evidence = {
            "uncategorized_file": file_upload.id,
            "access_activity_log": usage_markdown,
        }

        print(f"Evidence being submitted to stripe for dispute {dispute_id}:")
        # print(evidence)

        updated_dispute = stripe.Dispute.modify(
            dispute_id, evidence=evidence, submit=submit
        )

        if submit:
            logging.info(f"Dispute {dispute_id} SUBMITTED with file evidence")
        else:
            logging.info(f"Dispute {dispute_id} updated with file evidence")
        return updated_dispute

    except FileNotFoundError:
        logging.error(f"File not found: {file_path}")
        raise
    except InvalidRequestError as e:
        logging.error(f"Error updating dispute: {str(e)}")
        raise
    except APIError as e:
        logging.error(f"Stripe API error: {str(e)}")
        raise


def submit_dispute(dispute_id: str, output_path: str = None, submit: bool = False):
    dispute = get_dispute(dispute_id)
    dispute_info = hydrate_dispute_information(dispute)
    markdown_content = get_dispute_markdown(dispute_info)
    # print("=" * 100)
    # print("DISPUTE MARKDOWN:")
    # print("=" * 100)
    # print(markdown_content)

    # If no output path specified, create a default one based on dispute ID
    if not output_path:
        output_dir = "dispute_pdfs"
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        output_path = os.path.join(output_dir, f"dispute_{dispute_id}.pdf")

    create_and_save_pdf(markdown_content, output_path)
    update_dispute_details(
        dispute_id, dispute_info, file_path=output_path, submit=submit
    )


def process_disputes_from_file(
    file_path: str, output_dir: str = "dispute_pdfs", submit: bool = False
):
    """Process multiple disputes from a file containing one dispute ID per line using a thread pool."""
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"Created output directory: {output_dir}")

    with open(file_path, "r") as f:
        dispute_ids = [line.strip() for line in f if line.strip()]

    print(f"Found {len(dispute_ids)} dispute IDs to process")

    # Use a constant of 5 threads
    max_workers = 5
    print(f"Processing up to {max_workers} disputes concurrently")

    # Process disputes using a thread pool
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        # Submit all tasks and store the futures
        future_to_dispute = {
            executor.submit(
                process_single_dispute,
                dispute_id,
                output_dir,
                submit,
                i,
                len(dispute_ids),
            ): dispute_id
            for i, dispute_id in enumerate(dispute_ids, 1)
        }

        # Process results as they complete
        for future in concurrent.futures.as_completed(future_to_dispute):
            dispute_id = future_to_dispute[future]
            try:
                success, message = future.result()
                if success:
                    print(f"✅ {message}")
                else:
                    print(f"❌ {message}")
            except Exception as e:
                print(f"❌ Error processing dispute {dispute_id}: {str(e)}")


def process_single_dispute(
    dispute_id: str, output_dir: str, submit: bool, index: int, total: int
) -> Tuple[bool, str]:
    """Process a single dispute and return success status and message."""
    output_path = os.path.join(output_dir, f"dispute_{dispute_id}.pdf")
    try:
        print(f"Starting dispute: {dispute_id} ({index}/{total})")
        submit_dispute(dispute_id, output_path, submit)
        return (
            True,
            f"Successfully processed dispute {dispute_id} - {index}/{total} completed",
        )
    except Exception as e:
        return (
            False,
            f"Error processing dispute {dispute_id}: {str(e)} - {index}/{total} attempted",
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Fetch a Stripe dispute")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--dispute_id", help="ID of the dispute to retrieve")
    group.add_argument(
        "--dispute_ids_file", help="File containing one dispute ID per line"
    )

    parser.add_argument("--output", "-o", help="Output PDF file path or directory")
    parser.add_argument(
        "--execute",
        "-e",
        help="Execute the dispute submission, otherwise evidence is staged",
        action="store_true",
        default=False,
    )
    args = parser.parse_args()

    if args.dispute_ids_file:
        output_dir = args.output if args.output else "dispute_pdfs"
        process_disputes_from_file(
            args.dispute_ids_file, output_dir, submit=args.execute
        )
    else:
        submit_dispute(args.dispute_id, args.output, submit=args.execute)