"""Migrate Product prices

Revision ID: 44f8c56e3ddc
Revises: 7c77579439b1
Create Date: 2025-02-12 14:45:42.307332

"""

import sqlalchemy as sa
from alembic import op

# Polar Custom Imports

# revision identifiers, used by Alembic.
revision = "44f8c56e3ddc"
down_revision = "7c77579439b1"
branch_labels: tuple[str] | None = None
depends_on: tuple[str] | None = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.alter_column("product_prices", "type", existing_type=sa.VARCHAR(), nullable=True)
    op.add_column(
        "products", sa.Column("recurring_interval", sa.String(), nullable=True)
    )
    op.create_index(
        op.f("ix_products_recurring_interval"),
        "products",
        ["recurring_interval"],
        unique=False,
    )

    # One-time products: they have only one price
    # For those price, set their type to NULL
    op.execute(
        """
        UPDATE product_prices
        SET type = NULL
        WHERE type = 'one_time'
        """
    )

    # Recurring products that have only one price
    # Move the recurring interval to the product and set their type and recurring_interval to NULL
    op.execute(
        """
        WITH single_recurring_price_product AS (
            SELECT product_id
            FROM product_prices
            WHERE type = 'recurring' AND is_archived = FALSE
            GROUP BY product_id
            HAVING COUNT(*) = 1
        ), srp_product_price AS (
            SELECT products.id AS product_id, product_prices.id AS price_id, product_prices.recurring_interval AS recurring_interval
            FROM products
            JOIN product_prices ON products.id = product_prices.product_id
            WHERE products.id IN (SELECT product_id FROM single_recurring_price_product)
        ), product_update AS (
            UPDATE products
            SET recurring_interval = srp_product_price.recurring_interval
            FROM srp_product_price
            WHERE products.id = srp_product_price.product_id
        )
        UPDATE product_prices
        SET type = NULL, recurring_interval = NULL
        FROM srp_product_price
        WHERE product_prices.id = srp_product_price.price_id
        """
    )

    # Leave recurring products with multiple prices as they are

    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###

    op.execute(
        """
        WITH one_time_products AS (
            SELECT products.id
            FROM products
            WHERE recurring_interval IS NULL
        )
        UPDATE product_prices
        SET type = 'one_time'
        FROM one_time_products
        WHERE product_prices.product_id = one_time_products.id
        """
    )

    op.execute(
        """
        WITH recurring_products AS (
            SELECT products.id, products.recurring_interval
            FROM products
            WHERE recurring_interval IS NOT NULL
        )
        UPDATE product_prices
        SET type = 'recurring', recurring_interval = recurring_products.recurring_interval
        FROM recurring_products
        WHERE product_prices.product_id = recurring_products.id
        """
    )

    op.drop_index(op.f("ix_products_recurring_interval"), table_name="products")
    op.drop_column("products", "recurring_interval")
    op.alter_column(
        "product_prices", "type", existing_type=sa.VARCHAR(), nullable=False
    )
    # ### end Alembic commands ###
