"""Add Payout.paid_at

Revision ID: d21ac0812338
Revises: 2ffeb5016445
Create Date: 2025-06-19 15:15:28.126423

"""

import concurrent.futures
import random
import time
from datetime import UTC, datetime

import sqlalchemy as sa

# Polar Custom Imports
import stripe as stripe_lib
from alembic import op
from sqlalchemy.util.typing import TypedDict

from polar.config import settings

# revision identifiers, used by Alembic.
revision = "d21ac0812338"
down_revision = "2ffeb5016445"
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 MigratedPayout(TypedDict):
    payout_id: str
    paid_at: datetime


def migrate_payout(payout: tuple[str, str, str], retry: int = 1) -> MigratedPayout:
    payout_id, processor_id, account_stripe_id = payout
    try:
        payout_data = stripe_client.payouts.retrieve(
            processor_id,
            options={
                "stripe_account": account_stripe_id,
            },
        )
        return {
            "payout_id": payout_id,
            "paid_at": datetime.fromtimestamp(payout_data.arrival_date, UTC),
        }
    except stripe_lib.RateLimitError:
        time.sleep(retry + random.random())
        return migrate_payout(payout, retry=retry + 1)


def migrate_payouts(
    results: sa.CursorResult[tuple[str, str, str]],
) -> list[MigratedPayout]:
    migrated_payouts: list[MigratedPayout] = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
        futures = [
            executor.submit(migrate_payout, payout._tuple()) for payout in results
        ]
        for future in concurrent.futures.as_completed(futures):
            migrated_payouts.append(future.result())
    return migrated_payouts


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "payouts", sa.Column("paid_at", sa.TIMESTAMP(timezone=True), nullable=True)
    )

    connection = op.get_bind()
    payouts = connection.execute(
        sa.text("""
            SELECT payouts.id, payouts.processor_id, accounts.stripe_id
            FROM payouts
            JOIN accounts ON payouts.account_id = accounts.id
            WHERE payouts.processor_id IS NOT NULL
            AND payouts.paid_at IS NULL
            AND payouts.processor = 'stripe'
            AND payouts.status = 'succeeded'
        """)
    )
    migrated_payouts = migrate_payouts(payouts)
    for payout in migrated_payouts:
        connection.execute(
            sa.text("""
                UPDATE payouts
                SET paid_at = :paid_at
                WHERE id = :payout_id
            """).bindparams(
                payout_id=payout["payout_id"],
                paid_at=payout["paid_at"],
            )
        )

    # ### end Alembic commands ###


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