"""add embedding table and index

Revision ID: eee0433afdbf
Revises: ef63dfc830d3
Create Date: 2023-07-25 17:48:00.910876

"""

import pgvector
import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision = "eee0433afdbf"
down_revision = "ef63dfc830d3"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    connection = op.get_bind()
    result = connection.execute(
        sa.text("SELECT extname FROM pg_extension WHERE extname = 'vector';"),
    ).fetchone()

    if result:
        op.create_table(
            "embeddings",
            sa.Column("id", sa.String(), nullable=False),
            sa.Column("document_id", sa.String(), nullable=False),
            sa.Column("text", sa.String(), nullable=False),
            sa.Column("seq_num", sa.Integer(), nullable=False),
            sa.Column(
                "embedding",
                pgvector.sqlalchemy.Vector(dim=1536),
                nullable=False,
            ),
            sa.ForeignKeyConstraint(
                ["document_id"],
                ["documents.id"],
            ),
            sa.PrimaryKeyConstraint("id"),
        )
        op.create_index(
            "vector_index",
            "embeddings",
            ["embedding"],
            unique=False,
            postgresql_using="ivfflat",
            postgresql_with={"lists": 100},
            postgresql_ops={"embedding": "vector_l2_ops"},
        )
    else:
        op.create_table(
            "embeddings",
            sa.Column("id", sa.String(), nullable=False),
            sa.Column("document_id", sa.String(), nullable=False),
            sa.Column("text", sa.String(), nullable=False),
            sa.Column("seq_num", sa.Integer(), nullable=False),
            sa.Column("embedding", sa.String(), nullable=False),
            sa.ForeignKeyConstraint(
                ["document_id"],
                ["documents.id"],
            ),
            sa.PrimaryKeyConstraint("id"),
        )
    op.create_index(
        op.f("ix_embeddings_document_id"),
        "embeddings",
        ["document_id"],
        unique=False,
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    connection = op.get_bind()
    result = connection.execute(
        sa.text("SELECT extname FROM pg_extension WHERE extname = 'vector';"),
    ).fetchone()

    if result:
        op.drop_index(
            "my_index",
            table_name="embeddings",
            postgresql_using="ivfflat",
            postgresql_with={"lists": 100},
            postgresql_ops={"embedding": "vector_l2_ops"},
        )
    op.drop_index(op.f("ix_embeddings_document_id"), table_name="embeddings")
    op.drop_table("embeddings")
    # ### end Alembic commands ###
