"""add_is_admin_and_fix_connection_cascade

Revision ID: 9a977647c5a4
Revises: c60291fb2129
Create Date: 2025-10-06 17:30:05.584891

"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = '9a977647c5a4'
down_revision = '6e39730bc45b'
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint(op.f('connection_organization_id_fkey'), 'connection', type_='foreignkey')
    op.create_foreign_key(None, 'connection', 'organization', ['organization_id'], ['id'], ondelete='CASCADE')
    op.drop_index(op.f('idx_connection_init_session_redirect_session_id'), table_name='connection_init_session')
    op.alter_column('search_queries', 'collection_id',
               existing_type=sa.UUID(),
               comment='Collection that was searched',
               existing_nullable=False)
    op.alter_column('search_queries', 'user_id',
               existing_type=sa.UUID(),
               comment='User who performed the search (null for API key searches)',
               existing_nullable=True)
    op.alter_column('search_queries', 'api_key_id',
               existing_type=sa.UUID(),
               comment='API key used for the search (null for user searches)',
               existing_nullable=True)
    op.alter_column('search_queries', 'query_text',
               existing_type=sa.TEXT(),
               comment='The actual search query text',
               existing_nullable=False)
    op.alter_column('search_queries', 'query_length',
               existing_type=sa.INTEGER(),
               comment='Length of the search query in characters',
               existing_nullable=False)
    op.alter_column('search_queries', 'search_type',
               existing_type=sa.VARCHAR(length=20),
               comment="Type of search: 'basic', 'advanced'",
               existing_nullable=False)
    op.alter_column('search_queries', 'response_type',
               existing_type=sa.VARCHAR(length=20),
               comment="Response type: 'raw', 'completion'",
               existing_nullable=True)
    op.alter_column('search_queries', 'limit',
               existing_type=sa.INTEGER(),
               comment='Maximum number of results requested',
               existing_nullable=True)
    op.alter_column('search_queries', 'offset',
               existing_type=sa.INTEGER(),
               comment='Number of results to skip for pagination',
               existing_nullable=True)
    op.alter_column('search_queries', 'score_threshold',
               existing_type=sa.DOUBLE_PRECISION(precision=53),
               comment='Minimum similarity score threshold',
               existing_nullable=True)
    op.alter_column('search_queries', 'recency_bias',
               existing_type=sa.DOUBLE_PRECISION(precision=53),
               comment='Recency bias weight (0.0 to 1.0)',
               existing_nullable=True)
    op.alter_column('search_queries', 'search_method',
               existing_type=sa.VARCHAR(length=20),
               comment="Search method: 'hybrid', 'neural', 'keyword'",
               existing_nullable=True)
    op.alter_column('search_queries', 'duration_ms',
               existing_type=sa.INTEGER(),
               comment='Search execution time in milliseconds',
               existing_nullable=False)
    op.alter_column('search_queries', 'results_count',
               existing_type=sa.INTEGER(),
               comment='Number of results returned',
               existing_nullable=False)
    op.alter_column('search_queries', 'status',
               existing_type=sa.VARCHAR(length=20),
               comment="Search status: 'success', 'no_results', 'no_relevant_results', 'error'",
               existing_nullable=False)
    op.alter_column('search_queries', 'query_expansion_enabled',
               existing_type=sa.BOOLEAN(),
               comment='Whether query expansion was enabled',
               existing_nullable=True)
    op.alter_column('search_queries', 'reranking_enabled',
               existing_type=sa.BOOLEAN(),
               comment='Whether LLM reranking was enabled',
               existing_nullable=True)
    op.alter_column('search_queries', 'query_interpretation_enabled',
               existing_type=sa.BOOLEAN(),
               comment='Whether query interpretation was enabled',
               existing_nullable=True)

    # Add is_admin column as nullable first
    op.add_column('user', sa.Column('is_admin', sa.Boolean(), nullable=True))

    # Set default values based on email domain
    # First, set all users to is_admin = false
    op.execute("UPDATE \"user\" SET is_admin = false")

    # Then, set users with @airweave.ai email to is_admin = true
    op.execute("UPDATE \"user\" SET is_admin = true WHERE email LIKE '%@airweave.ai' OR email = 'admin@example.com'")

    # Now make the column NOT NULL
    op.alter_column('user', 'is_admin', nullable=False)
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('user', 'is_admin')
    op.alter_column('search_queries', 'query_interpretation_enabled',
               existing_type=sa.BOOLEAN(),
               comment=None,
               existing_comment='Whether query interpretation was enabled',
               existing_nullable=True)
    op.alter_column('search_queries', 'reranking_enabled',
               existing_type=sa.BOOLEAN(),
               comment=None,
               existing_comment='Whether LLM reranking was enabled',
               existing_nullable=True)
    op.alter_column('search_queries', 'query_expansion_enabled',
               existing_type=sa.BOOLEAN(),
               comment=None,
               existing_comment='Whether query expansion was enabled',
               existing_nullable=True)
    op.alter_column('search_queries', 'status',
               existing_type=sa.VARCHAR(length=20),
               comment=None,
               existing_comment="Search status: 'success', 'no_results', 'no_relevant_results', 'error'",
               existing_nullable=False)
    op.alter_column('search_queries', 'results_count',
               existing_type=sa.INTEGER(),
               comment=None,
               existing_comment='Number of results returned',
               existing_nullable=False)
    op.alter_column('search_queries', 'duration_ms',
               existing_type=sa.INTEGER(),
               comment=None,
               existing_comment='Search execution time in milliseconds',
               existing_nullable=False)
    op.alter_column('search_queries', 'search_method',
               existing_type=sa.VARCHAR(length=20),
               comment=None,
               existing_comment="Search method: 'hybrid', 'neural', 'keyword'",
               existing_nullable=True)
    op.alter_column('search_queries', 'recency_bias',
               existing_type=sa.DOUBLE_PRECISION(precision=53),
               comment=None,
               existing_comment='Recency bias weight (0.0 to 1.0)',
               existing_nullable=True)
    op.alter_column('search_queries', 'score_threshold',
               existing_type=sa.DOUBLE_PRECISION(precision=53),
               comment=None,
               existing_comment='Minimum similarity score threshold',
               existing_nullable=True)
    op.alter_column('search_queries', 'offset',
               existing_type=sa.INTEGER(),
               comment=None,
               existing_comment='Number of results to skip for pagination',
               existing_nullable=True)
    op.alter_column('search_queries', 'limit',
               existing_type=sa.INTEGER(),
               comment=None,
               existing_comment='Maximum number of results requested',
               existing_nullable=True)
    op.alter_column('search_queries', 'response_type',
               existing_type=sa.VARCHAR(length=20),
               comment=None,
               existing_comment="Response type: 'raw', 'completion'",
               existing_nullable=True)
    op.alter_column('search_queries', 'search_type',
               existing_type=sa.VARCHAR(length=20),
               comment=None,
               existing_comment="Type of search: 'basic', 'advanced'",
               existing_nullable=False)
    op.alter_column('search_queries', 'query_length',
               existing_type=sa.INTEGER(),
               comment=None,
               existing_comment='Length of the search query in characters',
               existing_nullable=False)
    op.alter_column('search_queries', 'query_text',
               existing_type=sa.TEXT(),
               comment=None,
               existing_comment='The actual search query text',
               existing_nullable=False)
    op.alter_column('search_queries', 'api_key_id',
               existing_type=sa.UUID(),
               comment=None,
               existing_comment='API key used for the search (null for user searches)',
               existing_nullable=True)
    op.alter_column('search_queries', 'user_id',
               existing_type=sa.UUID(),
               comment=None,
               existing_comment='User who performed the search (null for API key searches)',
               existing_nullable=True)
    op.alter_column('search_queries', 'collection_id',
               existing_type=sa.UUID(),
               comment=None,
               existing_comment='Collection that was searched',
               existing_nullable=False)
    op.create_index(op.f('idx_connection_init_session_redirect_session_id'), 'connection_init_session', ['redirect_session_id'], unique=False)
    op.drop_constraint(None, 'connection', type_='foreignkey')
    op.create_foreign_key(op.f('connection_organization_id_fkey'), 'connection', 'organization', ['organization_id'], ['id'])
    # ### end Alembic commands ###
