"""Remove Product.type/is_highlighted and introduce ProductPriceFree

Revision ID: 19f9bb88313b
Revises: 162775d9f4f2
Create Date: 2024-09-19 13:45:22.413448

"""

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

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

# Polar Custom Imports

# revision identifiers, used by Alembic.
revision = "19f9bb88313b"
down_revision = "162775d9f4f2"
branch_labels: tuple[str] | None = None
depends_on: tuple[str] | None = None


class MigratedProduct(TypedDict):
    product_id: UUID
    price_id: UUID
    stripe_product_id: str
    stripe_price_id: str


def process_product(
    product: tuple[UUID, str, str, UUID, str], retry: int = 1
) -> MigratedProduct:
    id, name, description, organization_id, organization_slug = product
    metadata = {"product_id": str(id)}
    metadata["organization_id"] = str(organization_id)
    metadata["organization_name"] = organization_slug

    product_idempotency_key = f"product-run2-{id}"
    try:
        stripe_product = stripe_lib.Product.create(
            name=f"{organization_slug} - {name}",
            description=description,
            metadata=metadata,
            default_price_data={
                "currency": "usd",
                "recurring": {"interval": "month"},
                "unit_amount": 0,
            },
            idempotency_key=product_idempotency_key,
        )
    except stripe_lib.RateLimitError:
        time.sleep(retry + random.random())
        return process_product(product, retry=retry + 1)

    price_id = uuid4()
    stripe_price_id = stripe_product.default_price
    assert isinstance(stripe_price_id, str)
    return {
        "product_id": id,
        "price_id": price_id,
        "stripe_product_id": stripe_product.id,
        "stripe_price_id": stripe_price_id,
    }


def migrate_products(results: sa.CursorResult[Any]) -> list[MigratedProduct]:
    migrated_products: list[MigratedProduct] = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
        future_to_product = {
            executor.submit(process_product, product._tuple()): product
            for product in results
        }
        for future in concurrent.futures.as_completed(future_to_product):
            migrated_products.append(future.result())
    return migrated_products


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    connection = op.get_bind()
    results = connection.execute(
        sa.text("""
            SELECT products.id, products.name, products.description, products.organization_id, organizations.slug
            FROM products
            LEFT JOIN organizations ON products.organization_id = organizations.id
            WHERE type = 'free'
        """)
    )
    migrated_products = migrate_products(results)

    op.alter_column(
        "product_prices",
        "price_currency",
        existing_type=sa.VARCHAR(length=3),
        nullable=True,
    )

    for migrated_product in migrated_products:
        op.execute(
            sa.text("""
            INSERT INTO product_prices (
                id,
                created_at,
                type,
                recurring_interval,
                amount_type,
                is_archived,
                stripe_price_id,
                product_id
            )
            VALUES (
                :id,
                NOW(),
                'recurring',
                'month',
                'free',
                FALSE,
                :stripe_price_id,
                :product_id
            )
            """).bindparams(
                id=migrated_product["price_id"],
                stripe_price_id=migrated_product["stripe_price_id"],
                product_id=migrated_product["product_id"],
            )
        )
        op.execute(
            sa.text(
                """
                UPDATE products
                SET stripe_product_id = :stripe_product_id
                WHERE id = :product_id
                """
            ).bindparams(
                stripe_product_id=migrated_product["stripe_product_id"],
                product_id=migrated_product["product_id"],
            )
        )
        op.execute(
            sa.text(
                """
                UPDATE subscriptions
                SET price_id = :price_id
                WHERE product_id = :product_id
                """
            ).bindparams(
                price_id=migrated_product["price_id"],
                product_id=migrated_product["product_id"],
            )
        )

    op.alter_column(
        "subscriptions",
        "price_id",
        existing_nullable=True,
        nullable=False,
    )
    op.drop_constraint(
        "subscriptions_price_id_fkey", "subscriptions", type_="foreignkey"
    )
    op.create_foreign_key(
        op.f("subscriptions_price_id_fkey"),
        "subscriptions",
        "product_prices",
        ["price_id"],
        ["id"],
        ondelete="cascade",
    )

    op.drop_index("idx_organization_id_type", table_name="products")
    op.drop_index("ix_products_is_highlighted", table_name="products")
    op.drop_index("ix_products_type", table_name="products")
    op.drop_column("products", "is_highlighted")
    op.drop_column("products", "type")
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "products", sa.Column("type", sa.VARCHAR(), autoincrement=False, nullable=True)
    )
    op.add_column(
        "products",
        sa.Column("is_highlighted", sa.BOOLEAN(), autoincrement=False, nullable=True),
    )
    op.create_index("ix_products_type", "products", ["type"], unique=False)
    op.create_index(
        "ix_products_is_highlighted", "products", ["is_highlighted"], unique=False
    )
    op.create_index(
        "idx_organization_id_type",
        "products",
        ["organization_id", "type"],
        unique=False,
    )

    op.alter_column(
        "subscriptions",
        "price_id",
        nullable=True,
    )
    op.drop_constraint(
        op.f("subscriptions_price_id_fkey"), "subscriptions", type_="foreignkey"
    )
    op.create_foreign_key(
        "subscriptions_price_id_fkey",
        "subscriptions",
        "product_prices",
        ["price_id"],
        ["id"],
        ondelete="SET NULL",
    )

    op.execute(
        """
        UPDATE products
        SET type = 'free', stripe_product_id = NULL
        WHERE id IN (
            SELECT product_id
            FROM product_prices
            WHERE amount_type = 'free' AND recurring_interval IS NOT NULL
        )
        """
    )
    op.execute(
        """
        DELETE FROM product_prices
        WHERE amount_type = 'free'
        """
    )

    op.alter_column(
        "product_prices", "stripe_price_id", existing_type=sa.VARCHAR(), nullable=False
    )
    op.alter_column(
        "product_prices",
        "price_currency",
        existing_type=sa.VARCHAR(length=3),
        nullable=False,
    )
    # ### end Alembic commands ###
