"""fix schema drift: sync indexes, constraints, and foreign key cascades

Revision ID: 87f8c5e11906
Revises: 33107394dcd1
Create Date: 2025-10-30 00:13:45.520219

"""

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision = "87f8c5e11906"
down_revision = "33107394dcd1"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.alter_column(
        "acquisition",
        "status",
        existing_type=postgresql.ENUM(
            "PLANNED", "STARTED", "COMPLETED", "PAUSED", "ABANDONED", name="acquisitionstatus"
        ),
        nullable=True,
        existing_server_default=sa.text("'PLANNED'::acquisitionstatus"),
    )

    # Drop FKs that depend on unique constraints first
    op.drop_constraint(
        op.f("agentinstructionacknowledgement_instruction_id_fkey"),
        "agentinstructionacknowledgement",
        type_="foreignkey",
    )
    op.drop_constraint(op.f("agentconnection_session_id_fkey"), "agentconnection", type_="foreignkey")
    op.drop_constraint(op.f("agentinstruction_session_id_fkey"), "agentinstruction", type_="foreignkey")

    # Now safe to modify agentconnection
    op.alter_column(
        "agentconnection",
        "id",
        existing_type=sa.BIGINT(),
        type_=sa.Integer(),
        existing_nullable=False,
        autoincrement=True,
    )
    op.drop_constraint(op.f("agentconnection_connection_id_key"), "agentconnection", type_="unique")
    op.drop_index(op.f("idx_agent_connection_agent_status"), table_name="agentconnection")
    op.create_index(op.f("ix_agentconnection_agent_id"), "agentconnection", ["agent_id"], unique=False)
    op.create_index(op.f("ix_agentconnection_connection_id"), "agentconnection", ["connection_id"], unique=True)
    op.create_index(op.f("ix_agentconnection_created_at"), "agentconnection", ["created_at"], unique=False)
    op.create_index(
        op.f("ix_agentconnection_last_heartbeat_at"), "agentconnection", ["last_heartbeat_at"], unique=False
    )
    op.create_index(op.f("ix_agentconnection_session_id"), "agentconnection", ["session_id"], unique=False)
    op.create_index(op.f("ix_agentconnection_status"), "agentconnection", ["status"], unique=False)

    # Now safe to modify agentinstruction
    op.alter_column(
        "agentinstruction",
        "id",
        existing_type=sa.BIGINT(),
        type_=sa.Integer(),
        existing_nullable=False,
        autoincrement=True,
        existing_server_default=sa.text("nextval('agentinstruction_id_seq'::regclass)"),
    )
    op.alter_column("agentinstruction", "payload", existing_type=postgresql.JSONB(astext_type=sa.Text()), nullable=True)
    op.drop_constraint(op.f("agentinstruction_instruction_id_key"), "agentinstruction", type_="unique")
    op.drop_index(op.f("idx_agent_instruction_agent_created"), table_name="agentinstruction")
    op.drop_index(op.f("idx_agent_instruction_session_status"), table_name="agentinstruction")
    op.create_index(op.f("ix_agentinstruction_acknowledged_at"), "agentinstruction", ["acknowledged_at"], unique=False)
    op.create_index(op.f("ix_agentinstruction_agent_id"), "agentinstruction", ["agent_id"], unique=False)
    op.create_index(op.f("ix_agentinstruction_created_at"), "agentinstruction", ["created_at"], unique=False)
    op.create_index(op.f("ix_agentinstruction_expires_at"), "agentinstruction", ["expires_at"], unique=False)
    op.create_index(op.f("ix_agentinstruction_instruction_id"), "agentinstruction", ["instruction_id"], unique=True)
    op.create_index(
        op.f("ix_agentinstruction_instruction_type"), "agentinstruction", ["instruction_type"], unique=False
    )
    op.create_index(op.f("ix_agentinstruction_priority"), "agentinstruction", ["priority"], unique=False)
    op.create_index(op.f("ix_agentinstruction_sent_at"), "agentinstruction", ["sent_at"], unique=False)
    op.create_index(op.f("ix_agentinstruction_sequence_number"), "agentinstruction", ["sequence_number"], unique=False)
    op.create_index(op.f("ix_agentinstruction_session_id"), "agentinstruction", ["session_id"], unique=False)
    op.create_index(op.f("ix_agentinstruction_status"), "agentinstruction", ["status"], unique=False)

    # Now safe to modify agentinstructionacknowledgement
    op.create_index(
        op.f("ix_agentinstructionacknowledgement_agent_id"),
        "agentinstructionacknowledgement",
        ["agent_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_agentinstructionacknowledgement_created_at"),
        "agentinstructionacknowledgement",
        ["created_at"],
        unique=False,
    )
    op.create_index(
        op.f("ix_agentinstructionacknowledgement_instruction_id"),
        "agentinstructionacknowledgement",
        ["instruction_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_agentinstructionacknowledgement_processed_at"),
        "agentinstructionacknowledgement",
        ["processed_at"],
        unique=False,
    )
    op.create_index(
        op.f("ix_agentinstructionacknowledgement_session_id"),
        "agentinstructionacknowledgement",
        ["session_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_agentinstructionacknowledgement_status"), "agentinstructionacknowledgement", ["status"], unique=False
    )

    # Now safe to modify agentsession
    op.alter_column(
        "agentsession",
        "id",
        existing_type=sa.BIGINT(),
        type_=sa.Integer(),
        existing_nullable=False,
        autoincrement=True,
        existing_server_default=sa.text("nextval('agentsession_id_seq'::regclass)"),
    )
    op.drop_constraint(op.f("agentsession_session_id_key"), "agentsession", type_="unique")
    op.drop_index(op.f("idx_agent_session_agent_status"), table_name="agentsession")
    op.create_index(op.f("ix_agentsession_acquisition_uuid"), "agentsession", ["acquisition_uuid"], unique=False)
    op.create_index(op.f("ix_agentsession_agent_id"), "agentsession", ["agent_id"], unique=False)
    op.create_index(op.f("ix_agentsession_created_at"), "agentsession", ["created_at"], unique=False)
    op.create_index(op.f("ix_agentsession_last_activity_at"), "agentsession", ["last_activity_at"], unique=False)
    op.create_index(op.f("ix_agentsession_session_id"), "agentsession", ["session_id"], unique=True)
    op.create_index(op.f("ix_agentsession_status"), "agentsession", ["status"], unique=False)
    op.drop_constraint(op.f("agentsession_acquisition_uuid_fkey"), "agentsession", type_="foreignkey")
    op.create_foreign_key(None, "agentsession", "acquisition", ["acquisition_uuid"], ["uuid"])
    op.drop_constraint(op.f("atlas_grid_uuid_fkey"), "atlas", type_="foreignkey")
    op.create_foreign_key(None, "atlas", "grid", ["grid_uuid"], ["uuid"])
    op.drop_constraint(op.f("atlastile_atlas_uuid_fkey"), "atlastile", type_="foreignkey")
    op.create_foreign_key(None, "atlastile", "atlas", ["atlas_uuid"], ["uuid"])
    op.drop_constraint(
        op.f("atlastilegridsquareposition_atlastile_uuid_fkey"), "atlastilegridsquareposition", type_="foreignkey"
    )
    op.drop_constraint(
        op.f("atlastilegridsquareposition_gridsquare_uuid_fkey"), "atlastilegridsquareposition", type_="foreignkey"
    )
    op.create_foreign_key(None, "atlastilegridsquareposition", "gridsquare", ["gridsquare_uuid"], ["uuid"])
    op.create_foreign_key(None, "atlastilegridsquareposition", "atlastile", ["atlastile_uuid"], ["uuid"])
    op.alter_column(
        "foilhole",
        "status",
        existing_type=postgresql.ENUM("NONE", "MICROGRAPHS_DETECTED", name="foilholestatus"),
        nullable=True,
        existing_server_default=sa.text("'NONE'::foilholestatus"),
    )
    op.drop_constraint(op.f("foilhole_gridsquare_uuid_fkey"), "foilhole", type_="foreignkey")
    op.create_foreign_key(None, "foilhole", "gridsquare", ["gridsquare_uuid"], ["uuid"])
    op.alter_column(
        "grid",
        "status",
        existing_type=postgresql.ENUM(
            "NONE",
            "SCAN_STARTED",
            "SCAN_COMPLETED",
            "GRID_SQUARES_DECISION_STARTED",
            "GRID_SQUARES_DECISION_COMPLETED",
            name="gridstatus",
        ),
        nullable=True,
        existing_server_default=sa.text("'NONE'::gridstatus"),
    )
    op.drop_constraint(op.f("grid_acquisition_uuid_fkey"), "grid", type_="foreignkey")
    op.create_foreign_key(None, "grid", "acquisition", ["acquisition_uuid"], ["uuid"])
    op.alter_column(
        "gridsquare",
        "status",
        existing_type=postgresql.ENUM(
            "NONE", "FOIL_HOLES_DECISION_STARTED", "FOIL_HOLES_DECISION_COMPLETED", name="gridsquarestatus"
        ),
        nullable=True,
        existing_server_default=sa.text("'NONE'::gridsquarestatus"),
    )
    op.drop_index(op.f("idx_gridsquare_acquisition_datetime_status"), table_name="gridsquare")
    op.drop_constraint(op.f("gridsquare_grid_uuid_fkey"), "gridsquare", type_="foreignkey")
    op.create_foreign_key(None, "gridsquare", "grid", ["grid_uuid"], ["uuid"])
    op.alter_column(
        "micrograph",
        "status",
        existing_type=postgresql.ENUM(
            "NONE",
            "MOTION_CORRECTION_STARTED",
            "MOTION_CORRECTION_COMPLETED",
            "CTF_STARTED",
            "CTF_COMPLETED",
            "PARTICLE_PICKING_STARTED",
            "PARTICLE_PICKING_COMPLETED",
            "PARTICLE_SELECTION_STARTED",
            "PARTICLE_SELECTION_COMPLETED",
            name="micrographstatus",
        ),
        nullable=True,
        existing_server_default=sa.text("'NONE'::micrographstatus"),
    )
    op.drop_index(op.f("idx_micrograph_acquisition_datetime"), table_name="micrograph")
    op.drop_constraint(op.f("micrograph_foilhole_uuid_fkey"), "micrograph", type_="foreignkey")
    op.create_foreign_key(None, "micrograph", "foilhole", ["foilhole_uuid"], ["uuid"])
    op.drop_constraint(op.f("qualityprediction_foilhole_uuid_fkey"), "qualityprediction", type_="foreignkey")
    op.drop_constraint(op.f("qualityprediction_gridsquare_uuid_fkey"), "qualityprediction", type_="foreignkey")
    op.drop_constraint(op.f("qualityprediction_prediction_model_name_fkey"), "qualityprediction", type_="foreignkey")
    op.create_foreign_key(None, "qualityprediction", "gridsquare", ["gridsquare_uuid"], ["uuid"])
    op.create_foreign_key(None, "qualityprediction", "foilhole", ["foilhole_uuid"], ["uuid"])
    op.create_foreign_key(None, "qualityprediction", "qualitypredictionmodel", ["prediction_model_name"], ["name"])
    op.drop_constraint(
        op.f("qualitypredictionmodelparameter_grid_uuid_fkey"), "qualitypredictionmodelparameter", type_="foreignkey"
    )
    op.drop_constraint(
        op.f("qualitypredictionmodelparameter_prediction_model_name_fkey"),
        "qualitypredictionmodelparameter",
        type_="foreignkey",
    )
    op.create_foreign_key(
        None, "qualitypredictionmodelparameter", "qualitypredictionmodel", ["prediction_model_name"], ["name"]
    )
    op.create_foreign_key(None, "qualitypredictionmodelparameter", "grid", ["grid_uuid"], ["uuid"])
    op.drop_constraint(
        op.f("qualitypredictionmodelweight_micrograph_uuid_fkey"), "qualitypredictionmodelweight", type_="foreignkey"
    )
    op.drop_constraint(
        op.f("qualitypredictionmodelweight_prediction_model_name_fkey"),
        "qualitypredictionmodelweight",
        type_="foreignkey",
    )
    op.drop_constraint(
        op.f("qualitypredictionmodelweight_grid_uuid_fkey"), "qualitypredictionmodelweight", type_="foreignkey"
    )
    op.create_foreign_key(
        None, "qualitypredictionmodelweight", "qualitypredictionmodel", ["prediction_model_name"], ["name"]
    )
    op.create_foreign_key(None, "qualitypredictionmodelweight", "micrograph", ["micrograph_uuid"], ["uuid"])
    op.create_foreign_key(None, "qualitypredictionmodelweight", "grid", ["grid_uuid"], ["uuid"])

    # Recreate FKs that we dropped at the beginning
    op.create_foreign_key(None, "agentconnection", "agentsession", ["session_id"], ["session_id"])
    op.create_foreign_key(None, "agentinstruction", "agentsession", ["session_id"], ["session_id"])
    op.create_foreign_key(
        None, "agentinstructionacknowledgement", "agentinstruction", ["instruction_id"], ["instruction_id"]
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint(None, "qualitypredictionmodelweight", type_="foreignkey")
    op.drop_constraint(None, "qualitypredictionmodelweight", type_="foreignkey")
    op.drop_constraint(None, "qualitypredictionmodelweight", type_="foreignkey")
    op.create_foreign_key(
        op.f("qualitypredictionmodelweight_grid_uuid_fkey"),
        "qualitypredictionmodelweight",
        "grid",
        ["grid_uuid"],
        ["uuid"],
        ondelete="CASCADE",
    )
    op.create_foreign_key(
        op.f("qualitypredictionmodelweight_prediction_model_name_fkey"),
        "qualitypredictionmodelweight",
        "qualitypredictionmodel",
        ["prediction_model_name"],
        ["name"],
        ondelete="CASCADE",
    )
    op.create_foreign_key(
        op.f("qualitypredictionmodelweight_micrograph_uuid_fkey"),
        "qualitypredictionmodelweight",
        "micrograph",
        ["micrograph_uuid"],
        ["uuid"],
        ondelete="CASCADE",
    )
    op.drop_constraint(None, "qualitypredictionmodelparameter", type_="foreignkey")
    op.drop_constraint(None, "qualitypredictionmodelparameter", type_="foreignkey")
    op.create_foreign_key(
        op.f("qualitypredictionmodelparameter_prediction_model_name_fkey"),
        "qualitypredictionmodelparameter",
        "qualitypredictionmodel",
        ["prediction_model_name"],
        ["name"],
        ondelete="CASCADE",
    )
    op.create_foreign_key(
        op.f("qualitypredictionmodelparameter_grid_uuid_fkey"),
        "qualitypredictionmodelparameter",
        "grid",
        ["grid_uuid"],
        ["uuid"],
        ondelete="CASCADE",
    )
    op.drop_constraint(None, "qualityprediction", type_="foreignkey")
    op.drop_constraint(None, "qualityprediction", type_="foreignkey")
    op.drop_constraint(None, "qualityprediction", type_="foreignkey")
    op.create_foreign_key(
        op.f("qualityprediction_prediction_model_name_fkey"),
        "qualityprediction",
        "qualitypredictionmodel",
        ["prediction_model_name"],
        ["name"],
        ondelete="CASCADE",
    )
    op.create_foreign_key(
        op.f("qualityprediction_gridsquare_uuid_fkey"),
        "qualityprediction",
        "gridsquare",
        ["gridsquare_uuid"],
        ["uuid"],
        ondelete="CASCADE",
    )
    op.create_foreign_key(
        op.f("qualityprediction_foilhole_uuid_fkey"),
        "qualityprediction",
        "foilhole",
        ["foilhole_uuid"],
        ["uuid"],
        ondelete="CASCADE",
    )
    op.drop_constraint(None, "micrograph", type_="foreignkey")
    op.create_foreign_key(
        op.f("micrograph_foilhole_uuid_fkey"), "micrograph", "foilhole", ["foilhole_uuid"], ["uuid"], ondelete="CASCADE"
    )
    op.create_index(op.f("idx_micrograph_acquisition_datetime"), "micrograph", ["acquisition_datetime"], unique=False)
    op.alter_column(
        "micrograph",
        "status",
        existing_type=postgresql.ENUM(
            "NONE",
            "MOTION_CORRECTION_STARTED",
            "MOTION_CORRECTION_COMPLETED",
            "CTF_STARTED",
            "CTF_COMPLETED",
            "PARTICLE_PICKING_STARTED",
            "PARTICLE_PICKING_COMPLETED",
            "PARTICLE_SELECTION_STARTED",
            "PARTICLE_SELECTION_COMPLETED",
            name="micrographstatus",
        ),
        nullable=False,
        existing_server_default=sa.text("'NONE'::micrographstatus"),
    )
    op.drop_constraint(None, "gridsquare", type_="foreignkey")
    op.create_foreign_key(
        op.f("gridsquare_grid_uuid_fkey"), "gridsquare", "grid", ["grid_uuid"], ["uuid"], ondelete="CASCADE"
    )
    op.create_index(
        op.f("idx_gridsquare_acquisition_datetime_status"),
        "gridsquare",
        ["acquisition_datetime", "status"],
        unique=False,
    )
    op.alter_column(
        "gridsquare",
        "status",
        existing_type=postgresql.ENUM(
            "NONE", "FOIL_HOLES_DECISION_STARTED", "FOIL_HOLES_DECISION_COMPLETED", name="gridsquarestatus"
        ),
        nullable=False,
        existing_server_default=sa.text("'NONE'::gridsquarestatus"),
    )
    op.drop_constraint(None, "grid", type_="foreignkey")
    op.create_foreign_key(
        op.f("grid_acquisition_uuid_fkey"), "grid", "acquisition", ["acquisition_uuid"], ["uuid"], ondelete="CASCADE"
    )
    op.alter_column(
        "grid",
        "status",
        existing_type=postgresql.ENUM(
            "NONE",
            "SCAN_STARTED",
            "SCAN_COMPLETED",
            "GRID_SQUARES_DECISION_STARTED",
            "GRID_SQUARES_DECISION_COMPLETED",
            name="gridstatus",
        ),
        nullable=False,
        existing_server_default=sa.text("'NONE'::gridstatus"),
    )
    op.drop_constraint(None, "foilhole", type_="foreignkey")
    op.create_foreign_key(
        op.f("foilhole_gridsquare_uuid_fkey"),
        "foilhole",
        "gridsquare",
        ["gridsquare_uuid"],
        ["uuid"],
        ondelete="CASCADE",
    )
    op.alter_column(
        "foilhole",
        "status",
        existing_type=postgresql.ENUM("NONE", "MICROGRAPHS_DETECTED", name="foilholestatus"),
        nullable=False,
        existing_server_default=sa.text("'NONE'::foilholestatus"),
    )
    op.drop_constraint(None, "atlastilegridsquareposition", type_="foreignkey")
    op.drop_constraint(None, "atlastilegridsquareposition", type_="foreignkey")
    op.create_foreign_key(
        op.f("atlastilegridsquareposition_gridsquare_uuid_fkey"),
        "atlastilegridsquareposition",
        "gridsquare",
        ["gridsquare_uuid"],
        ["uuid"],
        ondelete="CASCADE",
    )
    op.create_foreign_key(
        op.f("atlastilegridsquareposition_atlastile_uuid_fkey"),
        "atlastilegridsquareposition",
        "atlastile",
        ["atlastile_uuid"],
        ["uuid"],
        ondelete="CASCADE",
    )
    op.drop_constraint(None, "atlastile", type_="foreignkey")
    op.create_foreign_key(
        op.f("atlastile_atlas_uuid_fkey"), "atlastile", "atlas", ["atlas_uuid"], ["uuid"], ondelete="CASCADE"
    )
    op.drop_constraint(None, "atlas", type_="foreignkey")
    op.create_foreign_key(op.f("atlas_grid_uuid_fkey"), "atlas", "grid", ["grid_uuid"], ["uuid"], ondelete="CASCADE")
    op.drop_constraint(None, "agentsession", type_="foreignkey")
    op.create_foreign_key(
        op.f("agentsession_acquisition_uuid_fkey"),
        "agentsession",
        "acquisition",
        ["acquisition_uuid"],
        ["uuid"],
        ondelete="CASCADE",
    )
    op.drop_index(op.f("ix_agentsession_status"), table_name="agentsession")
    op.drop_index(op.f("ix_agentsession_session_id"), table_name="agentsession")
    op.drop_index(op.f("ix_agentsession_last_activity_at"), table_name="agentsession")
    op.drop_index(op.f("ix_agentsession_created_at"), table_name="agentsession")
    op.drop_index(op.f("ix_agentsession_agent_id"), table_name="agentsession")
    op.drop_index(op.f("ix_agentsession_acquisition_uuid"), table_name="agentsession")
    op.create_index(op.f("idx_agent_session_agent_status"), "agentsession", ["agent_id", "status"], unique=False)
    op.create_unique_constraint(op.f("agentsession_session_id_key"), "agentsession", ["session_id"])
    op.alter_column(
        "agentsession",
        "id",
        existing_type=sa.Integer(),
        type_=sa.BIGINT(),
        existing_nullable=False,
        autoincrement=True,
        existing_server_default=sa.text("nextval('agentsession_id_seq'::regclass)"),
    )
    op.drop_constraint(None, "agentinstructionacknowledgement", type_="foreignkey")
    op.drop_index(op.f("ix_agentinstructionacknowledgement_status"), table_name="agentinstructionacknowledgement")
    op.drop_index(op.f("ix_agentinstructionacknowledgement_session_id"), table_name="agentinstructionacknowledgement")
    op.drop_index(op.f("ix_agentinstructionacknowledgement_processed_at"), table_name="agentinstructionacknowledgement")
    op.drop_index(
        op.f("ix_agentinstructionacknowledgement_instruction_id"), table_name="agentinstructionacknowledgement"
    )
    op.drop_index(op.f("ix_agentinstructionacknowledgement_created_at"), table_name="agentinstructionacknowledgement")
    op.drop_index(op.f("ix_agentinstructionacknowledgement_agent_id"), table_name="agentinstructionacknowledgement")
    op.drop_constraint(None, "agentinstruction", type_="foreignkey")
    op.create_foreign_key(
        op.f("agentinstruction_session_id_fkey"),
        "agentinstruction",
        "agentsession",
        ["session_id"],
        ["session_id"],
        ondelete="CASCADE",
    )
    op.drop_index(op.f("ix_agentinstruction_status"), table_name="agentinstruction")
    op.drop_index(op.f("ix_agentinstruction_session_id"), table_name="agentinstruction")
    op.drop_index(op.f("ix_agentinstruction_sequence_number"), table_name="agentinstruction")
    op.drop_index(op.f("ix_agentinstruction_sent_at"), table_name="agentinstruction")
    op.drop_index(op.f("ix_agentinstruction_priority"), table_name="agentinstruction")
    op.drop_index(op.f("ix_agentinstruction_instruction_type"), table_name="agentinstruction")
    op.drop_index(op.f("ix_agentinstruction_instruction_id"), table_name="agentinstruction")
    op.drop_index(op.f("ix_agentinstruction_expires_at"), table_name="agentinstruction")
    op.drop_index(op.f("ix_agentinstruction_created_at"), table_name="agentinstruction")
    op.drop_index(op.f("ix_agentinstruction_agent_id"), table_name="agentinstruction")
    op.drop_index(op.f("ix_agentinstruction_acknowledged_at"), table_name="agentinstruction")
    op.create_index(
        op.f("idx_agent_instruction_session_status"), "agentinstruction", ["session_id", "status"], unique=False
    )
    op.create_index(
        op.f("idx_agent_instruction_agent_created"), "agentinstruction", ["agent_id", "created_at"], unique=False
    )
    op.create_unique_constraint(op.f("agentinstruction_instruction_id_key"), "agentinstruction", ["instruction_id"])
    op.create_foreign_key(
        op.f("agentinstructionacknowledgement_instruction_id_fkey"),
        "agentinstructionacknowledgement",
        "agentinstruction",
        ["instruction_id"],
        ["instruction_id"],
        ondelete="CASCADE",
    )
    op.alter_column(
        "agentinstruction", "payload", existing_type=postgresql.JSONB(astext_type=sa.Text()), nullable=False
    )
    op.alter_column(
        "agentinstruction",
        "id",
        existing_type=sa.Integer(),
        type_=sa.BIGINT(),
        existing_nullable=False,
        autoincrement=True,
        existing_server_default=sa.text("nextval('agentinstruction_id_seq'::regclass)"),
    )
    op.drop_constraint(None, "agentconnection", type_="foreignkey")
    op.create_foreign_key(
        op.f("agentconnection_session_id_fkey"),
        "agentconnection",
        "agentsession",
        ["session_id"],
        ["session_id"],
        ondelete="CASCADE",
    )
    op.drop_index(op.f("ix_agentconnection_status"), table_name="agentconnection")
    op.drop_index(op.f("ix_agentconnection_session_id"), table_name="agentconnection")
    op.drop_index(op.f("ix_agentconnection_last_heartbeat_at"), table_name="agentconnection")
    op.drop_index(op.f("ix_agentconnection_created_at"), table_name="agentconnection")
    op.drop_index(op.f("ix_agentconnection_connection_id"), table_name="agentconnection")
    op.drop_index(op.f("ix_agentconnection_agent_id"), table_name="agentconnection")
    op.create_index(op.f("idx_agent_connection_agent_status"), "agentconnection", ["agent_id", "status"], unique=False)
    op.create_unique_constraint(op.f("agentconnection_connection_id_key"), "agentconnection", ["connection_id"])
    op.alter_column(
        "agentconnection",
        "id",
        existing_type=sa.Integer(),
        type_=sa.BIGINT(),
        existing_nullable=False,
        autoincrement=True,
    )
    op.alter_column(
        "acquisition",
        "status",
        existing_type=postgresql.ENUM(
            "PLANNED", "STARTED", "COMPLETED", "PAUSED", "ABANDONED", name="acquisitionstatus"
        ),
        nullable=False,
        existing_server_default=sa.text("'PLANNED'::acquisitionstatus"),
    )
    # ### end Alembic commands ###
