"""add short_id to customers

Revision ID: aee75623c4bb
Revises: 8c9ef8060e95
Create Date: 2025-10-24 11:39:54.946342

"""

import sqlalchemy as sa
from alembic import op
from alembic_utils.pg_function import PGFunction
from sqlalchemy.schema import CreateSequence

# Polar Custom Imports
from polar.models import Customer

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


def upgrade() -> None:
    op.execute("CREATE SEQUENCE customer_short_id_seq START 1;")
    public_generate_customer_short_id = PGFunction(
        schema="public",
        signature="generate_customer_short_id(creation_timestamp timestamp with time zone DEFAULT clock_timestamp())",
        definition="returns bigint\n LANGUAGE plpgsql\nAS $function$\nDECLARE\n    our_epoch bigint := 1672531200000; -- 2023-01-01 in milliseconds\n    seq_id bigint;\n    now_millis bigint;\n    result bigint;\nBEGIN\n    -- Get sequence number modulo 1024 (10 bits)\n    SELECT nextval('customer_short_id_seq') % 1024 INTO seq_id;\n\n    -- Use provided timestamp (defaults to clock_timestamp())\n    SELECT FLOOR(EXTRACT(EPOCH FROM creation_timestamp) * 1000) INTO now_millis;\n\n    -- 42 bits timestamp (milliseconds) | 10 bits sequence = 52 bits total\n    -- Capacity\\: 1,024 IDs per millisecond (over 1 million per second)\n    -- Combine\\: (timestamp - epoch) << 10 | sequence\n    result := (now_millis - our_epoch) << 10;\n    result := result | seq_id;\n\n    RETURN result;\nEND;\n$function$",
    )
    op.create_entity(public_generate_customer_short_id)

    op.add_column(
        "customers",
        sa.Column("short_id", sa.BigInteger(), nullable=True),
    )

    op.execute(
        """
        UPDATE customers
        SET short_id = generate_customer_short_id(created_at);
        """
    )

    op.alter_column("customers", "short_id", nullable=False)

    op.execute(
        """
        ALTER TABLE customers
        ALTER COLUMN short_id SET DEFAULT generate_customer_short_id();
        """
    )

    op.create_unique_constraint(
        op.f("customers_organization_id_short_id_key"),
        "customers",
        ["organization_id", "short_id"],
    )
    op.create_index(
        op.f("ix_customers_short_id"), "customers", ["short_id"], unique=False
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(op.f("ix_customers_short_id"), table_name="customers")
    op.drop_constraint(
        op.f("customers_organization_id_short_id_key"), "customers", type_="unique"
    )
    op.drop_column("customers", "short_id")
    op.execute("DROP FUNCTION IF EXISTS generate_customer_short_id();")
    op.execute("DROP SEQUENCE IF EXISTS customer_short_id_seq;")
    # ### end Alembic commands ###
