"""Fix files_agents constraints

Revision ID: 1af251a42c06
Revises: 51999513bcf1
Create Date: 2025-06-30 11:50:42.200885

"""

from typing import Sequence, Union

from alembic import op
from letta.settings import settings

# revision identifiers, used by Alembic.
revision: str = "1af251a42c06"
down_revision: Union[str, None] = "51999513bcf1"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # Skip this migration for SQLite
    if not settings.letta_pg_uri_no_default:
        return

    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index("ix_files_agents_agent_file_name", table_name="files_agents")
    op.drop_index("ix_files_agents_file_id_agent_id", table_name="files_agents")
    op.drop_constraint("uq_files_agents_agent_file_name", "files_agents", type_="unique")
    op.drop_constraint("uq_files_agents_file_agent", "files_agents", type_="unique")
    op.create_index("ix_agent_filename", "files_agents", ["agent_id", "file_name"], unique=False)
    op.create_index("ix_file_agent", "files_agents", ["file_id", "agent_id"], unique=False)
    op.create_unique_constraint("uq_agent_filename", "files_agents", ["agent_id", "file_name"])
    op.create_unique_constraint("uq_file_agent", "files_agents", ["file_id", "agent_id"])
    # ### end Alembic commands ###


def downgrade() -> None:
    # Skip this migration for SQLite
    if not settings.letta_pg_uri_no_default:
        return

    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint("uq_file_agent", "files_agents", type_="unique")
    op.drop_constraint("uq_agent_filename", "files_agents", type_="unique")
    op.drop_index("ix_file_agent", table_name="files_agents")
    op.drop_index("ix_agent_filename", table_name="files_agents")
    op.create_unique_constraint("uq_files_agents_file_agent", "files_agents", ["file_id", "agent_id"], postgresql_nulls_not_distinct=False)
    op.create_unique_constraint(
        "uq_files_agents_agent_file_name", "files_agents", ["agent_id", "file_name"], postgresql_nulls_not_distinct=False
    )
    op.create_index("ix_files_agents_file_id_agent_id", "files_agents", ["file_id", "agent_id"], unique=False)
    op.create_index("ix_files_agents_agent_file_name", "files_agents", ["agent_id", "file_name"], unique=False)
    # ### end Alembic commands ###