"""Add Order.billing_reason

Revision ID: c4cf133e597c
Revises: 4f890422b67a
Create Date: 2024-10-15 09:19:20.839229

"""

import concurrent.futures
import random
import time
from typing import Any, TypedDict
from uuid import UUID

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

# Polar Custom Imports
from polar.config import settings

# revision identifiers, used by Alembic.
revision = "c4cf133e597c"
down_revision = "4f890422b67a"
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 OrderBillingReason(TypedDict):
    order_id: UUID
    billing_reason: str


def process_order(order: tuple[UUID, str], retry: int = 1) -> OrderBillingReason:
    order_id, stripe_invoice_id = order
    try:
        invoice = stripe_client.invoices.retrieve(stripe_invoice_id)
    except stripe_lib.RateLimitError:
        time.sleep(retry + random.random())
        return process_order(order, retry=retry + 1)

    if invoice.billing_reason not in [
        "subscription_create",
        "subscription_cycle",
        "subscription_update",
    ]:
        raise ValueError(f"Unexpected billing reason: {invoice.billing_reason}")

    return {
        "order_id": order_id,
        "billing_reason": invoice.billing_reason,
    }


def migrate_orders(results: sa.CursorResult[Any]) -> list[OrderBillingReason]:
    migrated_orders: list[OrderBillingReason] = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
        future_to_order = {
            executor.submit(process_order, order._tuple()): order for order in results
        }
        for future in concurrent.futures.as_completed(future_to_order):
            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_reason", sa.String(), nullable=True))

    op.execute(
        "UPDATE orders SET billing_reason = 'purchase' WHERE subscription_id IS NULL"
    )

    connection = op.get_bind()
    subscription_orders = connection.execute(
        sa.text("""
            SELECT orders.id, orders.stripe_invoice_id
            FROM orders
            WHERE orders.subscription_id IS NOT NULL
            AND orders.stripe_invoice_id IS NOT NULL
        """)
    )
    migrated_orders = migrate_orders(subscription_orders)
    for migrated_order in migrated_orders:
        op.execute(
            sa.text(
                """
                UPDATE orders
                SET billing_reason = :billing_reason
                WHERE id = :order_id
                """
            ).bindparams(
                order_id=migrated_order["order_id"],
                billing_reason=migrated_order["billing_reason"],
            )
        )

    op.alter_column("orders", "billing_reason", nullable=False)

    op.create_index(
        op.f("ix_orders_billing_reason"), "orders", ["billing_reason"], unique=False
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(op.f("ix_orders_billing_reason"), table_name="orders")
    op.drop_column("orders", "billing_reason")
    # ### end Alembic commands ###
