"""Add Order tax fields

Revision ID: d50c07506383
Revises: bba168067ce7
Create Date: 2025-05-23 09:34:52.985729

"""

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import postgresql

from polar.kit.tax import TaxIDType

# Polar Custom Imports

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


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

    op.execute(
        """
        UPDATE orders
        SET taxability_reason = 'standard_rated'
        WHERE tax_amount > 0
        """
    )
    op.execute(
        """
        UPDATE orders
        SET taxability_reason = 'reverse_charge'
        WHERE orders.customer_id IN (
            SELECT id FROM customers
            WHERE customers.tax_id IS NOT NULL
            AND customers.tax_id != 'null'::jsonb
        )
        """
    )
    op.execute(
        """
        UPDATE orders
        SET taxability_reason = 'not_collecting'
        WHERE taxability_reason IS NULL
        """
    )
    op.alter_column(
        "orders",
        "taxability_reason",
        existing_type=sa.String(),
        nullable=False,
        existing_nullable=True,
    )

    op.add_column(
        "orders",
        sa.Column("tax_id", postgresql.JSONB(astext_type=sa.Text()), nullable=True),
    )
    op.add_column(
        "orders",
        sa.Column("tax_rate", postgresql.JSONB(astext_type=sa.Text()), nullable=True),
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column("orders", "tax_rate")
    op.drop_column("orders", "tax_id")
    op.drop_column("orders", "taxability_reason")
    # ### end Alembic commands ###
