"""Add image generation fields to generation_records

Revision ID: 4b83fd481c7e
Revises: 88c3ab75b86e
Create Date: 2026-05-19 10:50:30.461713
"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '4b83fd481c7e'
down_revision: Union[str, None] = '88c3ab75b86e'
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! ###
    # First add columns with nullable=True for existing data
    op.add_column('generation_records', sa.Column('gen_type', sa.String(length=16), nullable=True))
    op.add_column('generation_records', sa.Column('image_size', sa.String(length=8), nullable=True))
    op.add_column('generation_records', sa.Column('image_url', sa.String(length=512), nullable=True))
    op.add_column('generation_records', sa.Column('image_tokens_used', sa.Integer(), nullable=True))
    
    # Update existing records with default values
    op.execute("UPDATE generation_records SET gen_type = 'video' WHERE gen_type IS NULL")
    op.execute("UPDATE generation_records SET image_tokens_used = 0 WHERE image_tokens_used IS NULL")
    
    # Now set columns to NOT NULL
    op.alter_column('generation_records', 'gen_type', nullable=False)
    op.alter_column('generation_records', 'image_tokens_used', nullable=False)
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('generation_records', 'image_tokens_used')
    op.drop_column('generation_records', 'image_url')
    op.drop_column('generation_records', 'image_size')
    op.drop_column('generation_records', 'gen_type')
    # ### end Alembic commands ###
