"""Add Order.billing_address

Revision ID: 1769a6e618a4
Revises: 6cbeabf73caf
Create Date: 2024-11-26 14:44:03.569035

"""

import concurrent.futures
import random
import time
from typing import Any, TypedDict, cast

import sqlalchemy as sa
import stripe as stripe_lib
from alembic import op
from pydantic import ValidationError

from polar.config import settings

# Polar Custom Imports
from polar.integrations.stripe.utils import get_expandable_id
from polar.kit.address import Address, AddressType

# revision identifiers, used by Alembic.
revision = "1769a6e618a4"
down_revision = "6cbeabf73caf"
branch_labels: tuple[str] | None = None
depends_on: tuple[str] | None = None


stripe_client = stripe_lib.StripeClient(
    settings.STRIPE_SECRET_KEY,
    http_client=stripe_lib.HTTPXClient(allow_sync_methods=True),
)


class MigratedOrder(TypedDict):
    order_id: str
    amount: int
    billing_address: dict[str, Any] | None


def _is_empty_customer_address(customer_address: dict[str, Any] | None) -> bool:
    return customer_address is None or customer_address["country"] is None


def migrate_order(
    order: tuple[str, int, str | None, str | None], retry: int = 1
) -> MigratedOrder:
    order_id, amount, stripe_invoice_id, stripe_charge_id = order

    if stripe_invoice_id is None and stripe_charge_id is None:
        raise ValueError(f"No invoice or charge: {order_id}")

    customer_address: Any | None = None
    try:
        # Get from invoice
        if stripe_invoice_id is not None:
            invoice = stripe_client.invoices.retrieve(stripe_invoice_id)
            customer_address = invoice.customer_address
            # No address on invoice, try to get from charge
            if (
                _is_empty_customer_address(customer_address)
                and invoice.charge is not None
            ):
                return migrate_order(
                    (order_id, amount, None, get_expandable_id(invoice.charge))
                )
        # Get from charge
        elif stripe_charge_id is not None:
            charge = stripe_client.charges.retrieve(
                stripe_charge_id,
                params={
                    "expand": ["payment_method_details", "payment_method_details.card"]
                },
            )
            customer_address = charge.billing_details.address
            # No address on charge, try to get from payment method
            if _is_empty_customer_address(customer_address):
                if payment_method_details := charge.payment_method_details:
                    if card := getattr(payment_method_details, "card", None):
                        customer_address = {"country": card.country}
    except stripe_lib.RateLimitError:
        time.sleep(retry + random.random())
        return migrate_order(order, retry=retry + 1)

    billing_address: dict[str, Any] | None = None
    if not _is_empty_customer_address(customer_address):
        try:
            billing_address = (
                Address.model_validate(customer_address).model_dump()
                if customer_address
                else None
            )
        except ValidationError as e:
            raise ValueError(f"Invalid address for order {order_id}: {e}")

    return {"order_id": order_id, "amount": amount, "billing_address": billing_address}


def migrate_orders(
    results: sa.CursorResult[tuple[str, int, str | None, str | None]],
) -> list[MigratedOrder]:
    migrated_orders: list[MigratedOrder] = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
        futures = [executor.submit(migrate_order, order._tuple()) for order in results]
        for future in concurrent.futures.as_completed(futures):
            migrated_orders.append(future.result())
    return migrated_orders


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "orders",
        sa.Column(
            "billing_address",
            AddressType(astext_type=sa.Text()),
            nullable=True,
        ),
    )

    connection = op.get_bind()
    orders = connection.execute(
        sa.text("""
            SELECT orders.id, orders.amount, orders.stripe_invoice_id, orders.user_metadata->>'charge_id' AS stripe_charge_id
            FROM orders
        """)
    )
    migrated_orders = migrate_orders(orders)
    for migrated_order in migrated_orders:
        if migrated_order["billing_address"] is None:
            if migrated_order["amount"] != 0:
                print("No billing address for paid order", migrated_order["order_id"])  # noqa: T201
            continue
        op.execute(
            sa.text(
                """
                UPDATE orders
                SET billing_address = :billing_address
                WHERE id = :order_id
                """
            ).bindparams(
                sa.bindparam(
                    "billing_address", migrated_order["billing_address"], type_=sa.JSON
                ),
                order_id=migrated_order["order_id"],
            )
        )

    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column("orders", "billing_address")
    # ### end Alembic commands ###
