"""add_tenant_chat_token_quota

Revision ID: 540c970ff890
Revises: 7fc8dd61bb95
Create Date: 2026-05-13 11:21:14.192033

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '540c970ff890'
down_revision: Union[str, None] = '7fc8dd61bb95'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # Get current connection and inspector
    conn = op.get_bind()
    inspector = sa.inspect(conn)
    
    # Check if columns already exist (defensive migration)
    columns = [c['name'] for c in inspector.get_columns('tenants')]
    
    if 'chat_token_daily_limit' not in columns:
        op.add_column('tenants', sa.Column('chat_token_daily_limit', sa.BigInteger(), nullable=False, server_default='1000000'))
    
    if 'chat_tokens_used_today' not in columns:
        op.add_column('tenants', sa.Column('chat_tokens_used_today', sa.BigInteger(), nullable=False, server_default='0'))
        
    if 'chat_tokens_last_reset_at' not in columns:
        op.add_column('tenants', sa.Column('chat_tokens_last_reset_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False))


def downgrade() -> None:
    op.drop_column('tenants', 'chat_tokens_last_reset_at')
    op.drop_column('tenants', 'chat_tokens_used_today')
    op.drop_column('tenants', 'chat_token_daily_limit')