"""added quota tables

Revision ID: 2159d4ac431a
Revises: e0d1a9b9f6c4
Create Date: 2025-07-08 14:38:38.542486

"""

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision: str = "2159d4ac431a"
down_revision: Union[str, None] = "e0d1a9b9f6c4"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    # 1) Create the `quota_type` enum *before* we add the column that references it.
    #    Alembic does **not** automatically emit a CREATE TYPE statement for Enum
    #    objects when they are used inside `op.add_column()`. If we skip this step
    #    Postgres raises `type "quota_type" does not exist` when applying the
    #    migration.
    quota_type_enum = sa.Enum("monthly", "annual", name="quota_type")
    quota_type_enum.create(op.get_bind(), checkfirst=True)

    op.create_table(
        "organization_usage_cycles",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("organization_id", sa.Integer(), nullable=False),
        sa.Column("period_start", sa.DateTime(), nullable=False),
        sa.Column("period_end", sa.DateTime(), nullable=False),
        sa.Column("quota_dograh_tokens", sa.Integer(), nullable=False),
        sa.Column("used_dograh_tokens", sa.Integer(), nullable=False),
        sa.Column("created_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("updated_at", sa.DateTime(timezone=True), nullable=True),
        sa.ForeignKeyConstraint(
            ["organization_id"],
            ["organizations.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint(
            "organization_id", "period_start", "period_end", name="unique_org_period"
        ),
    )
    op.create_index(
        "idx_usage_cycles_org_period",
        "organization_usage_cycles",
        ["organization_id", "period_end"],
        unique=False,
    )
    op.create_index(
        op.f("ix_organization_usage_cycles_id"),
        "organization_usage_cycles",
        ["id"],
        unique=False,
    )
    # Add the column now that the enum type exists.
    op.add_column(
        "organizations",
        sa.Column(
            "quota_type", quota_type_enum, nullable=False, server_default="monthly"
        ),
    )
    op.add_column(
        "organizations",
        sa.Column(
            "quota_dograh_tokens",
            sa.Integer(),
            nullable=False,
            server_default=sa.text("0"),
        ),
    )
    op.add_column(
        "organizations",
        sa.Column(
            "quota_reset_day",
            sa.Integer(),
            nullable=False,
            server_default=sa.text("LEAST(EXTRACT(DAY FROM CURRENT_DATE)::int, 28)"),
        ),
    )
    op.add_column(
        "organizations", sa.Column("quota_start_date", sa.DateTime(), nullable=True)
    )
    op.add_column(
        "organizations",
        sa.Column(
            "quota_enabled",
            sa.Boolean(),
            nullable=False,
            server_default=sa.text("false"),
        ),
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column("organizations", "quota_enabled")
    op.drop_column("organizations", "quota_start_date")
    op.drop_column("organizations", "quota_reset_day")
    op.drop_column("organizations", "quota_dograh_tokens")
    op.drop_column("organizations", "quota_type")
    op.drop_index(
        op.f("ix_organization_usage_cycles_id"), table_name="organization_usage_cycles"
    )
    op.drop_index("idx_usage_cycles_org_period", table_name="organization_usage_cycles")
    op.drop_table("organization_usage_cycles")

    # Drop the enum type now that no column depends on it.
    quota_type_enum = sa.Enum(name="quota_type")
    quota_type_enum.drop(op.get_bind(), checkfirst=True)
    # ### end Alembic commands ###