"""add rag prompts evaluation status conversation model dataset lineage

Revision ID: 245c81353f35
Revises: b94d8782dbef
Create Date: 2026-07-23 22:40:42.218663

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
import pgvector.sqlalchemy


# revision identifiers, used by Alembic.
revision: str = '245c81353f35'
down_revision: Union[str, Sequence[str], None] = 'b94d8782dbef'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    """Upgrade schema."""
    op.execute('CREATE EXTENSION IF NOT EXISTS vector')
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('prompts',
    sa.Column('org_id', sa.UUID(), nullable=False),
    sa.Column('name', sa.String(length=255), nullable=False),
    sa.Column('domain', sa.String(length=100), nullable=True),
    sa.Column('description', sa.Text(), nullable=True),
    sa.Column('content', sa.Text(), nullable=False),
    sa.Column('current_version', sa.Integer(), nullable=False),
    sa.Column('is_active', sa.Boolean(), nullable=False),
    sa.Column('created_by', sa.UUID(), nullable=True),
    sa.Column('id', sa.UUID(), nullable=False),
    sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False),
    sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False),
    sa.ForeignKeyConstraint(['created_by'], ['users.id'], ),
    sa.ForeignKeyConstraint(['org_id'], ['organizations.id'], ),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_table('prompt_versions',
    sa.Column('prompt_id', sa.UUID(), nullable=False),
    sa.Column('version', sa.Integer(), nullable=False),
    sa.Column('content', sa.Text(), nullable=False),
    sa.Column('created_by', sa.UUID(), nullable=True),
    sa.Column('id', sa.UUID(), nullable=False),
    sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False),
    sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False),
    sa.ForeignKeyConstraint(['created_by'], ['users.id'], ),
    sa.ForeignKeyConstraint(['prompt_id'], ['prompts.id'], ),
    sa.PrimaryKeyConstraint('id')
    )
    op.add_column('api_keys', sa.Column('last_used_at', sa.DateTime(timezone=True), nullable=True))
    op.alter_column('api_keys', 'revoked_at',
               existing_type=sa.VARCHAR(length=50),
               type_=sa.DateTime(timezone=True),
               existing_nullable=True,
               postgresql_using='revoked_at::timestamptz')
    op.add_column('conversations', sa.Column('model_name', sa.String(length=255), nullable=True))
    op.add_column('dataset_examples', sa.Column('parent_example_id', sa.UUID(), nullable=True))
    op.add_column('dataset_examples', sa.Column('source_conversation_id', sa.UUID(), nullable=True))
    op.create_foreign_key(None, 'dataset_examples', 'conversations', ['source_conversation_id'], ['id'])
    op.create_foreign_key(None, 'dataset_examples', 'dataset_examples', ['parent_example_id'], ['id'])
    # evaluation_results has no rows yet in this deployment (module was previously a
    # stub), so project_id can go straight to NOT NULL with no backfill.
    op.add_column('evaluation_results', sa.Column('project_id', sa.UUID(), nullable=False))
    op.add_column('evaluation_results', sa.Column('base_model', sa.String(length=255), nullable=True))
    evaluation_status_enum = sa.Enum('queued', 'running', 'completed', 'failed', name='evaluation_status')
    evaluation_status_enum.create(op.get_bind(), checkfirst=True)
    op.add_column('evaluation_results', sa.Column('status', evaluation_status_enum, nullable=False, server_default='completed'))
    op.alter_column('evaluation_results', 'status', server_default=None)
    op.add_column('evaluation_results', sa.Column('created_by', sa.UUID(), nullable=True))
    op.alter_column('evaluation_results', 'model_version_id',
               existing_type=sa.UUID(),
               nullable=True)
    op.create_foreign_key(None, 'evaluation_results', 'projects', ['project_id'], ['id'])
    op.create_foreign_key(None, 'evaluation_results', 'users', ['created_by'], ['id'])
    op.add_column('vector_embeddings', sa.Column('chunk_index', sa.Integer(), nullable=False, server_default='0'))
    op.alter_column('vector_embeddings', 'chunk_index', server_default=None)
    op.add_column('vector_embeddings', sa.Column('embedding', pgvector.sqlalchemy.Vector(768), nullable=True))
    # ### end Alembic commands ###


def downgrade() -> None:
    """Downgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('vector_embeddings', 'embedding')
    op.drop_column('vector_embeddings', 'chunk_index')
    op.drop_constraint(None, 'evaluation_results', type_='foreignkey')
    op.drop_constraint(None, 'evaluation_results', type_='foreignkey')
    op.alter_column('evaluation_results', 'model_version_id',
               existing_type=sa.UUID(),
               nullable=False)
    op.drop_column('evaluation_results', 'created_by')
    op.drop_column('evaluation_results', 'status')
    op.drop_column('evaluation_results', 'base_model')
    op.drop_column('evaluation_results', 'project_id')
    op.drop_constraint(None, 'dataset_examples', type_='foreignkey')
    op.drop_constraint(None, 'dataset_examples', type_='foreignkey')
    op.drop_column('dataset_examples', 'source_conversation_id')
    op.drop_column('dataset_examples', 'parent_example_id')
    op.drop_column('conversations', 'model_name')
    op.alter_column('api_keys', 'revoked_at',
               existing_type=sa.DateTime(timezone=True),
               type_=sa.VARCHAR(length=50),
               existing_nullable=True)
    op.drop_column('api_keys', 'last_used_at')
    op.drop_table('prompt_versions')
    op.drop_table('prompts')
    # ### end Alembic commands ###