"""Migrate agents to orm

Revision ID: d05669b60ebe
Revises: c5d964280dff
Create Date: 2024-12-12 10:25:31.825635

"""

from typing import Sequence, Union

import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

from alembic import op
from letta.settings import settings

# revision identifiers, used by Alembic.
revision: str = "d05669b60ebe"
down_revision: Union[str, None] = "c5d964280dff"
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.create_table(
        "sources_agents",
        sa.Column("agent_id", sa.String(), nullable=False),
        sa.Column("source_id", sa.String(), nullable=False),
        sa.ForeignKeyConstraint(
            ["agent_id"],
            ["agents.id"],
        ),
        sa.ForeignKeyConstraint(
            ["source_id"],
            ["sources.id"],
        ),
        sa.PrimaryKeyConstraint("agent_id", "source_id"),
    )
    op.drop_index("agent_source_mapping_idx_user", table_name="agent_source_mapping")
    op.drop_table("agent_source_mapping")
    op.add_column("agents", sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=True))
    op.add_column("agents", sa.Column("is_deleted", sa.Boolean(), server_default=sa.text("FALSE"), nullable=False))
    op.add_column("agents", sa.Column("_created_by_id", sa.String(), nullable=True))
    op.add_column("agents", sa.Column("_last_updated_by_id", sa.String(), nullable=True))
    op.add_column("agents", sa.Column("organization_id", sa.String(), nullable=True))
    # Populate `organization_id` based on `user_id`
    # Use a raw SQL query to update the organization_id
    op.execute(
        """
        UPDATE agents
        SET organization_id = users.organization_id
        FROM users
        WHERE agents.user_id = users.id
    """
    )
    op.alter_column("agents", "organization_id", nullable=False)
    op.alter_column("agents", "name", existing_type=sa.VARCHAR(), nullable=True)
    op.drop_index("agents_idx_user", table_name="agents")
    op.create_unique_constraint("unique_org_agent_name", "agents", ["organization_id", "name"])
    op.create_foreign_key(None, "agents", "organizations", ["organization_id"], ["id"])
    op.drop_column("agents", "tool_names")
    op.drop_column("agents", "user_id")
    op.drop_constraint("agents_tags_organization_id_fkey", "agents_tags", type_="foreignkey")
    op.drop_column("agents_tags", "_created_by_id")
    op.drop_column("agents_tags", "_last_updated_by_id")
    op.drop_column("agents_tags", "updated_at")
    op.drop_column("agents_tags", "id")
    op.drop_column("agents_tags", "is_deleted")
    op.drop_column("agents_tags", "created_at")
    op.drop_column("agents_tags", "organization_id")
    op.create_unique_constraint("unique_agent_block", "blocks_agents", ["agent_id", "block_id"])
    op.drop_constraint("fk_block_id_label", "blocks_agents", type_="foreignkey")
    op.create_foreign_key(
        "fk_block_id_label", "blocks_agents", "block", ["block_id", "block_label"], ["id", "label"], initially="DEFERRED", deferrable=True
    )
    op.drop_column("blocks_agents", "_created_by_id")
    op.drop_column("blocks_agents", "_last_updated_by_id")
    op.drop_column("blocks_agents", "updated_at")
    op.drop_column("blocks_agents", "id")
    op.drop_column("blocks_agents", "is_deleted")
    op.drop_column("blocks_agents", "created_at")
    op.drop_constraint("unique_tool_per_agent", "tools_agents", type_="unique")
    op.create_unique_constraint("unique_agent_tool", "tools_agents", ["agent_id", "tool_id"])
    op.drop_constraint("fk_tool_id", "tools_agents", type_="foreignkey")
    op.drop_constraint("tools_agents_agent_id_fkey", "tools_agents", type_="foreignkey")
    op.create_foreign_key(None, "tools_agents", "tools", ["tool_id"], ["id"], ondelete="CASCADE")
    op.create_foreign_key(None, "tools_agents", "agents", ["agent_id"], ["id"], ondelete="CASCADE")
    op.drop_column("tools_agents", "_created_by_id")
    op.drop_column("tools_agents", "tool_name")
    op.drop_column("tools_agents", "_last_updated_by_id")
    op.drop_column("tools_agents", "updated_at")
    op.drop_column("tools_agents", "id")
    op.drop_column("tools_agents", "is_deleted")
    op.drop_column("tools_agents", "created_at")
    # ### 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.add_column(
        "tools_agents",
        sa.Column("created_at", postgresql.TIMESTAMP(timezone=True), server_default=sa.text("now()"), autoincrement=False, nullable=True),
    )
    op.add_column(
        "tools_agents", sa.Column("is_deleted", sa.BOOLEAN(), server_default=sa.text("false"), autoincrement=False, nullable=False)
    )
    op.add_column("tools_agents", sa.Column("id", sa.VARCHAR(), autoincrement=False, nullable=False))
    op.add_column(
        "tools_agents",
        sa.Column("updated_at", postgresql.TIMESTAMP(timezone=True), server_default=sa.text("now()"), autoincrement=False, nullable=True),
    )
    op.add_column("tools_agents", sa.Column("_last_updated_by_id", sa.VARCHAR(), autoincrement=False, nullable=True))
    op.add_column("tools_agents", sa.Column("tool_name", sa.VARCHAR(), autoincrement=False, nullable=False))
    op.add_column("tools_agents", sa.Column("_created_by_id", sa.VARCHAR(), autoincrement=False, nullable=True))
    op.drop_constraint(None, "tools_agents", type_="foreignkey")
    op.drop_constraint(None, "tools_agents", type_="foreignkey")
    op.create_foreign_key("tools_agents_agent_id_fkey", "tools_agents", "agents", ["agent_id"], ["id"])
    op.create_foreign_key("fk_tool_id", "tools_agents", "tools", ["tool_id"], ["id"])
    op.drop_constraint("unique_agent_tool", "tools_agents", type_="unique")
    op.create_unique_constraint("unique_tool_per_agent", "tools_agents", ["agent_id", "tool_name"])
    op.add_column(
        "blocks_agents",
        sa.Column("created_at", postgresql.TIMESTAMP(timezone=True), server_default=sa.text("now()"), autoincrement=False, nullable=True),
    )
    op.add_column(
        "blocks_agents", sa.Column("is_deleted", sa.BOOLEAN(), server_default=sa.text("false"), autoincrement=False, nullable=False)
    )
    op.add_column("blocks_agents", sa.Column("id", sa.VARCHAR(), autoincrement=False, nullable=False))
    op.add_column(
        "blocks_agents",
        sa.Column("updated_at", postgresql.TIMESTAMP(timezone=True), server_default=sa.text("now()"), autoincrement=False, nullable=True),
    )
    op.add_column("blocks_agents", sa.Column("_last_updated_by_id", sa.VARCHAR(), autoincrement=False, nullable=True))
    op.add_column("blocks_agents", sa.Column("_created_by_id", sa.VARCHAR(), autoincrement=False, nullable=True))
    op.drop_constraint("fk_block_id_label", "blocks_agents", type_="foreignkey")
    op.create_foreign_key("fk_block_id_label", "blocks_agents", "block", ["block_id", "block_label"], ["id", "label"])
    op.drop_constraint("unique_agent_block", "blocks_agents", type_="unique")
    op.add_column("agents_tags", sa.Column("organization_id", sa.VARCHAR(), autoincrement=False, nullable=False))
    op.add_column(
        "agents_tags",
        sa.Column("created_at", postgresql.TIMESTAMP(timezone=True), server_default=sa.text("now()"), autoincrement=False, nullable=True),
    )
    op.add_column(
        "agents_tags", sa.Column("is_deleted", sa.BOOLEAN(), server_default=sa.text("false"), autoincrement=False, nullable=False)
    )
    op.add_column("agents_tags", sa.Column("id", sa.VARCHAR(), autoincrement=False, nullable=False))
    op.add_column(
        "agents_tags",
        sa.Column("updated_at", postgresql.TIMESTAMP(timezone=True), server_default=sa.text("now()"), autoincrement=False, nullable=True),
    )
    op.add_column("agents_tags", sa.Column("_last_updated_by_id", sa.VARCHAR(), autoincrement=False, nullable=True))
    op.add_column("agents_tags", sa.Column("_created_by_id", sa.VARCHAR(), autoincrement=False, nullable=True))
    op.create_foreign_key("agents_tags_organization_id_fkey", "agents_tags", "organizations", ["organization_id"], ["id"])
    op.add_column("agents", sa.Column("user_id", sa.VARCHAR(), autoincrement=False, nullable=False))
    op.add_column("agents", sa.Column("tool_names", postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=True))
    op.drop_constraint(None, "agents", type_="foreignkey")
    op.drop_constraint("unique_org_agent_name", "agents", type_="unique")
    op.create_index("agents_idx_user", "agents", ["user_id"], unique=False)
    op.alter_column("agents", "name", existing_type=sa.VARCHAR(), nullable=False)
    op.drop_column("agents", "organization_id")
    op.drop_column("agents", "_last_updated_by_id")
    op.drop_column("agents", "_created_by_id")
    op.drop_column("agents", "is_deleted")
    op.drop_column("agents", "updated_at")
    op.create_table(
        "agent_source_mapping",
        sa.Column("id", sa.VARCHAR(), autoincrement=False, nullable=False),
        sa.Column("user_id", sa.VARCHAR(), autoincrement=False, nullable=False),
        sa.Column("agent_id", sa.VARCHAR(), autoincrement=False, nullable=False),
        sa.Column("source_id", sa.VARCHAR(), autoincrement=False, nullable=False),
        sa.PrimaryKeyConstraint("id", name="agent_source_mapping_pkey"),
    )
    op.create_index("agent_source_mapping_idx_user", "agent_source_mapping", ["user_id", "agent_id", "source_id"], unique=False)
    op.drop_table("sources_agents")
    # ### end Alembic commands ###