"""Add chat token usage tracking

Revision ID: cd27562afa2e
Revises: 1f4d2a8c9b77
Create Date: 2026-04-16 13:05:09.572442

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision: str = 'cd27562afa2e'
down_revision: Union[str, None] = '1f4d2a8c9b77'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def _table_exists(bind, table_name: str) -> bool:
    return sa.inspect(bind).has_table(table_name)


def upgrade() -> None:
    bind = op.get_bind()

    if _table_exists(bind, 'chatbot_documents'):
        with op.batch_alter_table('chatbot_documents', schema=None) as batch_op:
            batch_op.alter_column('chunks_count',
                   existing_type=sa.INTEGER(),
                   nullable=False)
            batch_op.alter_column('created_at',
                   existing_type=postgresql.TIMESTAMP(),
                   nullable=False)

    if _table_exists(bind, 'vector_chunks'):
        with op.batch_alter_table('vector_chunks', schema=None) as batch_op:
            batch_op.alter_column('document_id',
                   existing_type=sa.VARCHAR(),
                   nullable=False)
            batch_op.alter_column('chunk_text',
                   existing_type=sa.VARCHAR(),
                   type_=sa.Text(),
                   existing_nullable=False)
            batch_op.alter_column('chunk_index',
                   existing_type=sa.INTEGER(),
                   nullable=False)
            batch_op.alter_column('created_at',
                   existing_type=postgresql.TIMESTAMP(),
                   type_=sa.DateTime(timezone=True),
                   nullable=False)
            batch_op.drop_constraint('vector_chunks_user_id_fkey', type_='foreignkey')
            batch_op.create_foreign_key(None, 'users', ['user_id'], ['id'], ondelete='CASCADE')

    if _table_exists(bind, 'vector_indices'):
        with op.batch_alter_table('vector_indices', schema=None) as batch_op:
            batch_op.alter_column('total_vectors',
                   existing_type=sa.INTEGER(),
                   nullable=False)
            batch_op.alter_column('updated_at',
                   existing_type=postgresql.TIMESTAMP(),
                   type_=sa.DateTime(timezone=True),
                   nullable=False)
            batch_op.drop_constraint('vector_indices_user_id_fkey', type_='foreignkey')
            batch_op.create_foreign_key(None, 'users', ['user_id'], ['id'], ondelete='CASCADE')


def downgrade() -> None:
    bind = op.get_bind()

    if _table_exists(bind, 'vector_indices'):
        with op.batch_alter_table('vector_indices', schema=None) as batch_op:
            batch_op.drop_constraint(None, type_='foreignkey')
            batch_op.create_foreign_key('vector_indices_user_id_fkey', 'users', ['user_id'], ['id'])
            batch_op.drop_index(batch_op.f('ix_vector_indices_user_id'))
            batch_op.create_index('ix_vector_indices_user_id', ['user_id'], unique=True)
            batch_op.create_index('ix_vector_indices_id', ['id'], unique=False)
            batch_op.alter_column('updated_at',
                   existing_type=sa.DateTime(timezone=True),
                   type_=postgresql.TIMESTAMP(),
                   nullable=True)
            batch_op.alter_column('total_vectors',
                   existing_type=sa.INTEGER(),
                   nullable=True)

    if _table_exists(bind, 'vector_chunks'):
        with op.batch_alter_table('vector_chunks', schema=None) as batch_op:
            batch_op.drop_constraint(None, type_='foreignkey')
            batch_op.create_foreign_key('vector_chunks_user_id_fkey', 'users', ['user_id'], ['id'])
            batch_op.create_index('ix_vector_chunks_id', ['id'], unique=False)
            batch_op.alter_column('created_at',
                   existing_type=sa.DateTime(timezone=True),
                   type_=postgresql.TIMESTAMP(),
                   nullable=True)
            batch_op.alter_column('chunk_index',
                   existing_type=sa.INTEGER(),
                   nullable=True)
            batch_op.alter_column('chunk_text',
                   existing_type=sa.Text(),
                   type_=sa.VARCHAR(),
                   existing_nullable=False)
            batch_op.alter_column('document_id',
                   existing_type=sa.VARCHAR(),
                   nullable=True)

    if _table_exists(bind, 'chatbot_documents'):
        with op.batch_alter_table('chatbot_documents', schema=None) as batch_op:
            batch_op.create_index('ix_chatbot_documents_id', ['id'], unique=False)
            batch_op.create_index('ix_chatbot_documents_drive_file_id', ['drive_file_id'], unique=False)
            batch_op.alter_column('created_at',
                   existing_type=postgresql.TIMESTAMP(),
                   nullable=True)
            batch_op.alter_column('chunks_count',
                   existing_type=sa.INTEGER(),
                   nullable=True)