"""checkout_links label and product

Revision ID: 2996c0ed1512
Revises: ec0834c42223
Create Date: 2024-11-12 17:10:57.321303

"""

import sqlalchemy as sa
from alembic import op

# Polar Custom Imports

# revision identifiers, used by Alembic.
revision = "2996c0ed1512"
down_revision = "ec0834c42223"
branch_labels: tuple[str] | None = None
depends_on: tuple[str] | None = None


def upgrade_products_on_existing_links() -> None:
    op.execute(
        """
        WITH mapping AS (
            SELECT
                checkout_links.id,
                product_prices.product_id
            FROM checkout_links
            JOIN product_prices ON product_prices.id = checkout_links.product_price_id
        )
        UPDATE checkout_links
        SET product_id = mapping.product_id
        FROM mapping
        WHERE 1 = 1
            AND checkout_links.id = mapping.id
        """
    )


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column("checkout_links", sa.Column("label", sa.String(), nullable=True))

    # Allow NULL until we can set all historic records in migration
    op.add_column("checkout_links", sa.Column("product_id", sa.Uuid(), nullable=True))
    upgrade_products_on_existing_links()
    op.alter_column("checkout_links", "product_id", nullable=False)
    op.create_foreign_key(
        op.f("checkout_links_product_id_fkey"),
        "checkout_links",
        "products",
        ["product_id"],
        ["id"],
        ondelete="cascade",
    )

    op.alter_column("checkout_links", "product_price_id", nullable=True)
    # ### end Alembic commands ###


def downgrade_set_prices_on_existing_links() -> None:
    op.execute(
        """
        WITH mapping AS (
            SELECT
                checkout_links.id AS checkout_link_id,
                product_prices.id AS product_price_id,
                product_prices.created_at AS price_created_at,
                RANK() OVER (PARTITION BY checkout_links.id ORDER BY product_prices.created_at) AS price_order
            FROM checkout_links
            JOIN product_prices ON product_prices.product_id = checkout_links.product_id
            WHERE 1 = 1
                AND checkout_links.product_price_id IS NULL
                AND product_prices.is_archived = false
        )
        UPDATE checkout_links
        SET product_price_id = mapping.product_price_id
        FROM mapping
        WHERE 1 = 1
            AND checkout_links.id = mapping.checkout_link_id
            AND mapping.price_order = 1
        """
    )


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint(
        op.f("checkout_links_product_id_fkey"), "checkout_links", type_="foreignkey"
    )
    downgrade_set_prices_on_existing_links()
    op.alter_column(
        "checkout_links", "product_price_id", existing_type=sa.UUID(), nullable=False
    )
    op.drop_column("checkout_links", "product_id")
    op.drop_column("checkout_links", "label")
    # ### end Alembic commands ###
