"""Add vector_db_provider to archives table

Revision ID: 068588268b02
Revises: d5103ee17ed5
Create Date: 2025-08-27 13:16:29.428231

"""

from typing import Sequence, Union

import sqlalchemy as sa

from alembic import op
from letta.settings import settings

# revision identifiers, used by Alembic.
revision: str = "068588268b02"
down_revision: Union[str, None] = "887a4367b560"
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! ###
    if settings.letta_pg_uri_no_default:
        # PostgreSQL - use enum type
        vectordbprovider = sa.Enum("NATIVE", "TPUF", name="vectordbprovider")
        vectordbprovider.create(op.get_bind(), checkfirst=True)

        # Add column as nullable first
        op.add_column("archives", sa.Column("vector_db_provider", vectordbprovider, nullable=True))

        # Backfill existing rows with NATIVE
        op.execute("UPDATE archives SET vector_db_provider = 'NATIVE' WHERE vector_db_provider IS NULL")

        # Make column non-nullable
        op.alter_column("archives", "vector_db_provider", nullable=False)
    else:
        # SQLite - use string type
        # Add column as nullable first
        op.add_column("archives", sa.Column("vector_db_provider", sa.String(), nullable=True))

        # Backfill existing rows with NATIVE
        op.execute("UPDATE archives SET vector_db_provider = 'NATIVE' WHERE vector_db_provider IS NULL")

        # For SQLite, we need to recreate the table to make column non-nullable
        # This is a limitation of SQLite ALTER TABLE
        # For simplicity, we'll leave it nullable in SQLite
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column("archives", "vector_db_provider")

    if settings.letta_pg_uri_no_default:
        # Drop enum type for PostgreSQL
        vectordbprovider = sa.Enum("NATIVE", "TPUF", name="vectordbprovider")
        vectordbprovider.drop(op.get_bind(), checkfirst=True)
    # ### end Alembic commands ###
