"""add datasource,rag processing info to document

Revision ID: 090822101cb5
Revises: 7cfa662dff86
Create Date: 2024-07-03 04:28:31.842226

"""

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 = "090822101cb5"
down_revision: Union[str, None] = "7cfa662dff86"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None

document_status = postgresql.ENUM(
    "SUCCESS", "PROCESSING", "ERROR", name="document_status"
)


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    document_status.create(op.get_bind())

    op.add_column(
        "datasource", sa.Column("schedule", sa.Unicode(length=255), nullable=True)
    )

    op.add_column(
        "document", sa.Column("datasource_id", sa.Unicode(length=255), nullable=True)
    )
    op.add_column(
        "document",
        sa.Column(
            "status", sa.Enum("SUCCESS", "PROCESSING", "ERROR", name="document_status")
        ),
    )
    op.add_column("document", sa.Column("last_modified", sa.DateTime()))
    op.add_column("document", sa.Column("last_processed", sa.DateTime()))
    op.add_column(
        "document", sa.Column("last_processed_message", sa.Text(), nullable=True)
    )
    op.alter_column(
        "document",
        "base_uri",
        existing_type=sa.VARCHAR(length=255),
        type_=sa.Unicode(length=4096),
        existing_nullable=False,
    )
    op.alter_column(
        "document",
        "filename",
        existing_type=sa.VARCHAR(length=255),
        type_=sa.Unicode(length=4096),
        existing_nullable=False,
    )
    op.alter_column(
        "document",
        "rag_metadata",
        existing_type=postgresql.JSONB(astext_type=sa.Text()),
        nullable=True,
    )

    op.drop_constraint("uq_document_uuid_base_url_filename", "document", type_="unique")
    op.create_unique_constraint(
        "uq_document_uuid_datasource_id", "document", ["uuid", "datasource_id"]
    )

    op.create_index(op.f("idx_document_base_uri"), "document", ["base_uri"])
    # ### end Alembic commands ###


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

    op.alter_column(
        "document",
        "rag_metadata",
        existing_type=postgresql.JSONB(astext_type=sa.Text()),
        nullable=False,
    )
    op.alter_column(
        "document",
        "filename",
        existing_type=sa.Unicode(length=4096),
        type_=sa.VARCHAR(length=255),
        existing_nullable=False,
    )
    op.alter_column(
        "document",
        "base_uri",
        existing_type=sa.Unicode(length=4096),
        type_=sa.VARCHAR(length=255),
        existing_nullable=False,
    )
    op.drop_column("document", "last_processed_message")
    op.drop_column("document", "last_processed")
    op.drop_column("document", "last_modified")
    op.drop_column("document", "status")
    op.drop_column("document", "datasource_id")

    document_status.drop(op.get_bind())

    # op.drop_constraint("uq_document_uuid_datasource_id", "document", type_="unique")
    op.create_unique_constraint(
        "uq_document_uuid_base_url_filename",
        "document",
        ["uuid", "base_uri", "filename"],
    )

    op.drop_index(op.f("idx_document_base_uri"), table_name="document")

    # ### end Alembic commands ###
