"""add subscription and usage tables

Revision ID: 3ca3dee4b44a
Revises: 1c98dcbcf360
Create Date: 2026-04-01 19:34:05.628630
"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '3ca3dee4b44a'
down_revision: Union[str, None] = '1c98dcbcf360'
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! ###
    op.create_table('subscriptions',
    sa.Column('id', sa.UUID(), nullable=False),
    sa.Column('org_id', sa.UUID(), nullable=False),
    sa.Column('plan', sa.Enum('HOBBY', 'PRO', 'STARTUP', 'ENTERPRISE', name='subscription_plan', native_enum=False, length=20), nullable=False),
    sa.Column('status', sa.Enum('ACTIVE', 'PAST_DUE', 'CANCELED', 'INCOMPLETE', name='subscription_status', native_enum=False, length=20), nullable=False),
    sa.Column('stripe_customer_id', sa.String(length=255), nullable=True),
    sa.Column('stripe_subscription_id', sa.String(length=255), nullable=True),
    sa.Column('current_period_start', sa.DateTime(timezone=True), nullable=False),
    sa.Column('current_period_end', sa.DateTime(timezone=True), nullable=False),
    sa.Column('canceled_at', sa.DateTime(timezone=True), nullable=True),
    sa.Column('created_at', sa.DateTime(timezone=True), nullable=False),
    sa.Column('updated_at', sa.DateTime(timezone=True), nullable=False),
    sa.ForeignKeyConstraint(['org_id'], ['organizations.id'], ondelete='CASCADE'),
    sa.PrimaryKeyConstraint('id'),
    sa.UniqueConstraint('org_id'),
    sa.UniqueConstraint('stripe_subscription_id')
    )
    op.create_index('ix_subscriptions_org_id', 'subscriptions', ['org_id'], unique=False)
    op.create_index(op.f('ix_subscriptions_stripe_customer_id'), 'subscriptions', ['stripe_customer_id'], unique=False)
    op.create_table('usage_records',
    sa.Column('id', sa.UUID(), nullable=False),
    sa.Column('org_id', sa.UUID(), nullable=False),
    sa.Column('period_start', sa.DateTime(timezone=True), nullable=False),
    sa.Column('period_end', sa.DateTime(timezone=True), nullable=False),
    sa.Column('trace_count', sa.Integer(), nullable=False),
    sa.Column('trace_eval_count', sa.Integer(), nullable=False),
    sa.Column('session_eval_count', sa.Integer(), nullable=False),
    sa.Column('monitoring_request_count', sa.Integer(), nullable=False),
    sa.Column('billed', sa.Boolean(), nullable=False),
    sa.Column('stripe_invoice_id', sa.String(length=255), nullable=True),
    sa.Column('created_at', sa.DateTime(timezone=True), nullable=False),
    sa.Column('updated_at', sa.DateTime(timezone=True), nullable=False),
    sa.ForeignKeyConstraint(['org_id'], ['organizations.id'], ondelete='CASCADE'),
    sa.PrimaryKeyConstraint('id'),
    sa.UniqueConstraint('org_id', 'period_start', name='uq_usage_record_org_period')
    )
    op.create_index('ix_usage_records_org_period', 'usage_records', ['org_id', 'period_start'], unique=False)
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index('ix_usage_records_org_period', table_name='usage_records')
    op.drop_table('usage_records')
    op.drop_index(op.f('ix_subscriptions_stripe_customer_id'), table_name='subscriptions')
    op.drop_index('ix_subscriptions_org_id', table_name='subscriptions')
    op.drop_table('subscriptions')
    # ### end Alembic commands ###
