"""Add multi-org support and nullable user tracking

Revision ID: d16c82cfe1fe
Revises: b1170e606aae
Create Date: 2025-06-05 16:49:11.311598

"""

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

# revision identifiers, used by Alembic.
revision = "d16c82cfe1fe"
down_revision = "b1170e606aae"
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "user_organization",
        sa.Column("user_id", sa.UUID(), nullable=False),
        sa.Column("organization_id", sa.UUID(), nullable=False),
        sa.Column("auth0_org_id", sa.String(), nullable=True),
        sa.Column("role", sa.String(), nullable=False),
        sa.Column("is_primary", sa.Boolean(), nullable=False),
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("modified_at", sa.DateTime(), nullable=False),
        sa.ForeignKeyConstraint(
            ["organization_id"],
            ["organization.id"],
        ),
        sa.ForeignKeyConstraint(
            ["user_id"],
            ["user.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.drop_index(op.f("ix_chat_message_chat_id"), table_name="chat_message")
    op.drop_table("chat_message")
    op.drop_table("chat")
    op.alter_column("api_key", "created_by_email", existing_type=sa.VARCHAR(), nullable=True)
    op.alter_column("api_key", "modified_by_email", existing_type=sa.VARCHAR(), nullable=True)
    op.drop_constraint(op.f("api_key_created_by_email_fkey"), "api_key", type_="foreignkey")
    op.drop_constraint(op.f("api_key_modified_by_email_fkey"), "api_key", type_="foreignkey")
    op.alter_column("collection", "created_by_email", existing_type=sa.VARCHAR(), nullable=True)
    op.alter_column("collection", "modified_by_email", existing_type=sa.VARCHAR(), nullable=True)
    op.drop_constraint(op.f("collection_created_by_email_fkey"), "collection", type_="foreignkey")
    op.drop_constraint(op.f("collection_modified_by_email_fkey"), "collection", type_="foreignkey")
    op.alter_column("dag_edge", "created_by_email", existing_type=sa.VARCHAR(), nullable=True)
    op.alter_column("dag_edge", "modified_by_email", existing_type=sa.VARCHAR(), nullable=True)
    op.drop_constraint(op.f("dag_edge_created_by_email_fkey"), "dag_edge", type_="foreignkey")
    op.drop_constraint(op.f("dag_edge_modified_by_email_fkey"), "dag_edge", type_="foreignkey")
    op.alter_column("dag_node", "created_by_email", existing_type=sa.VARCHAR(), nullable=True)
    op.alter_column("dag_node", "modified_by_email", existing_type=sa.VARCHAR(), nullable=True)
    op.drop_constraint(op.f("dag_node_modified_by_email_fkey"), "dag_node", type_="foreignkey")
    op.drop_constraint(op.f("dag_node_created_by_email_fkey"), "dag_node", type_="foreignkey")
    op.alter_column(
        "integration_credential", "created_by_email", existing_type=sa.VARCHAR(), nullable=True
    )
    op.alter_column(
        "integration_credential", "modified_by_email", existing_type=sa.VARCHAR(), nullable=True
    )
    op.drop_constraint(
        op.f("integration_credential_created_by_email_fkey"),
        "integration_credential",
        type_="foreignkey",
    )
    op.drop_constraint(
        op.f("integration_credential_modified_by_email_fkey"),
        "integration_credential",
        type_="foreignkey",
    )
    op.add_column("organization", sa.Column("auth0_org_id", sa.String(), nullable=True))
    op.alter_column(
        "source_connection", "created_by_email", existing_type=sa.VARCHAR(), nullable=True
    )
    op.alter_column(
        "source_connection", "modified_by_email", existing_type=sa.VARCHAR(), nullable=True
    )
    op.drop_constraint(
        op.f("source_connection_modified_by_email_fkey"), "source_connection", type_="foreignkey"
    )
    op.drop_constraint(
        op.f("source_connection_created_by_email_fkey"), "source_connection", type_="foreignkey"
    )
    op.alter_column("sync", "created_by_email", existing_type=sa.VARCHAR(), nullable=True)
    op.alter_column("sync", "modified_by_email", existing_type=sa.VARCHAR(), nullable=True)
    op.drop_constraint(op.f("sync_created_by_email_fkey"), "sync", type_="foreignkey")
    op.drop_constraint(op.f("sync_modified_by_email_fkey"), "sync", type_="foreignkey")
    op.alter_column("sync_dag", "created_by_email", existing_type=sa.VARCHAR(), nullable=True)
    op.alter_column("sync_dag", "modified_by_email", existing_type=sa.VARCHAR(), nullable=True)
    op.drop_constraint(op.f("sync_dag_created_by_email_fkey"), "sync_dag", type_="foreignkey")
    op.drop_constraint(op.f("sync_dag_modified_by_email_fkey"), "sync_dag", type_="foreignkey")
    op.alter_column("sync_job", "created_by_email", existing_type=sa.VARCHAR(), nullable=True)
    op.alter_column("sync_job", "modified_by_email", existing_type=sa.VARCHAR(), nullable=True)
    op.drop_constraint(op.f("sync_job_created_by_email_fkey"), "sync_job", type_="foreignkey")
    op.drop_constraint(op.f("sync_job_modified_by_email_fkey"), "sync_job", type_="foreignkey")
    op.add_column("user", sa.Column("primary_organization_id", sa.UUID(), nullable=True))
    op.add_column("user", sa.Column("current_organization_id", sa.UUID(), nullable=True))
    op.create_foreign_key(None, "user", "organization", ["primary_organization_id"], ["id"])
    op.alter_column("white_label", "created_by_email", existing_type=sa.VARCHAR(), nullable=True)
    op.alter_column("white_label", "modified_by_email", existing_type=sa.VARCHAR(), nullable=True)
    op.drop_constraint(op.f("white_label_created_by_email_fkey"), "white_label", type_="foreignkey")
    op.drop_constraint(
        op.f("white_label_modified_by_email_fkey"), "white_label", type_="foreignkey"
    )
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_foreign_key(
        op.f("white_label_modified_by_email_fkey"),
        "white_label",
        "user",
        ["modified_by_email"],
        ["email"],
    )
    op.create_foreign_key(
        op.f("white_label_created_by_email_fkey"),
        "white_label",
        "user",
        ["created_by_email"],
        ["email"],
    )
    op.alter_column("white_label", "modified_by_email", existing_type=sa.VARCHAR(), nullable=False)
    op.alter_column("white_label", "created_by_email", existing_type=sa.VARCHAR(), nullable=False)
    op.drop_constraint(None, "user", type_="foreignkey")
    op.drop_column("user", "current_organization_id")
    op.drop_column("user", "primary_organization_id")
    op.create_foreign_key(
        op.f("sync_job_modified_by_email_fkey"),
        "sync_job",
        "user",
        ["modified_by_email"],
        ["email"],
    )
    op.create_foreign_key(
        op.f("sync_job_created_by_email_fkey"), "sync_job", "user", ["created_by_email"], ["email"]
    )
    op.alter_column("sync_job", "modified_by_email", existing_type=sa.VARCHAR(), nullable=False)
    op.alter_column("sync_job", "created_by_email", existing_type=sa.VARCHAR(), nullable=False)
    op.create_foreign_key(
        op.f("sync_dag_modified_by_email_fkey"),
        "sync_dag",
        "user",
        ["modified_by_email"],
        ["email"],
    )
    op.create_foreign_key(
        op.f("sync_dag_created_by_email_fkey"), "sync_dag", "user", ["created_by_email"], ["email"]
    )
    op.alter_column("sync_dag", "modified_by_email", existing_type=sa.VARCHAR(), nullable=False)
    op.alter_column("sync_dag", "created_by_email", existing_type=sa.VARCHAR(), nullable=False)
    op.create_foreign_key(
        op.f("sync_modified_by_email_fkey"), "sync", "user", ["modified_by_email"], ["email"]
    )
    op.create_foreign_key(
        op.f("sync_created_by_email_fkey"), "sync", "user", ["created_by_email"], ["email"]
    )
    op.alter_column("sync", "modified_by_email", existing_type=sa.VARCHAR(), nullable=False)
    op.alter_column("sync", "created_by_email", existing_type=sa.VARCHAR(), nullable=False)
    op.create_foreign_key(
        op.f("source_connection_created_by_email_fkey"),
        "source_connection",
        "user",
        ["created_by_email"],
        ["email"],
    )
    op.create_foreign_key(
        op.f("source_connection_modified_by_email_fkey"),
        "source_connection",
        "user",
        ["modified_by_email"],
        ["email"],
    )
    op.alter_column(
        "source_connection", "modified_by_email", existing_type=sa.VARCHAR(), nullable=False
    )
    op.alter_column(
        "source_connection", "created_by_email", existing_type=sa.VARCHAR(), nullable=False
    )
    op.drop_column("organization", "auth0_org_id")
    op.create_foreign_key(
        op.f("integration_credential_modified_by_email_fkey"),
        "integration_credential",
        "user",
        ["modified_by_email"],
        ["email"],
    )
    op.create_foreign_key(
        op.f("integration_credential_created_by_email_fkey"),
        "integration_credential",
        "user",
        ["created_by_email"],
        ["email"],
    )
    op.alter_column(
        "integration_credential", "modified_by_email", existing_type=sa.VARCHAR(), nullable=False
    )
    op.alter_column(
        "integration_credential", "created_by_email", existing_type=sa.VARCHAR(), nullable=False
    )
    op.create_foreign_key(
        op.f("dag_node_created_by_email_fkey"), "dag_node", "user", ["created_by_email"], ["email"]
    )
    op.create_foreign_key(
        op.f("dag_node_modified_by_email_fkey"),
        "dag_node",
        "user",
        ["modified_by_email"],
        ["email"],
    )
    op.alter_column("dag_node", "modified_by_email", existing_type=sa.VARCHAR(), nullable=False)
    op.alter_column("dag_node", "created_by_email", existing_type=sa.VARCHAR(), nullable=False)
    op.create_foreign_key(
        op.f("dag_edge_modified_by_email_fkey"),
        "dag_edge",
        "user",
        ["modified_by_email"],
        ["email"],
    )
    op.create_foreign_key(
        op.f("dag_edge_created_by_email_fkey"), "dag_edge", "user", ["created_by_email"], ["email"]
    )
    op.alter_column("dag_edge", "modified_by_email", existing_type=sa.VARCHAR(), nullable=False)
    op.alter_column("dag_edge", "created_by_email", existing_type=sa.VARCHAR(), nullable=False)
    op.create_foreign_key(
        op.f("collection_modified_by_email_fkey"),
        "collection",
        "user",
        ["modified_by_email"],
        ["email"],
    )
    op.create_foreign_key(
        op.f("collection_created_by_email_fkey"),
        "collection",
        "user",
        ["created_by_email"],
        ["email"],
    )
    op.alter_column("collection", "modified_by_email", existing_type=sa.VARCHAR(), nullable=False)
    op.alter_column("collection", "created_by_email", existing_type=sa.VARCHAR(), nullable=False)
    op.create_foreign_key(
        op.f("api_key_modified_by_email_fkey"), "api_key", "user", ["modified_by_email"], ["email"]
    )
    op.create_foreign_key(
        op.f("api_key_created_by_email_fkey"), "api_key", "user", ["created_by_email"], ["email"]
    )
    op.alter_column("api_key", "modified_by_email", existing_type=sa.VARCHAR(), nullable=False)
    op.alter_column("api_key", "created_by_email", existing_type=sa.VARCHAR(), nullable=False)
    op.create_table(
        "chat_message",
        sa.Column("chat_id", sa.UUID(), autoincrement=False, nullable=False),
        sa.Column("content", sa.TEXT(), autoincrement=False, nullable=False),
        sa.Column(
            "role",
            postgresql.ENUM("USER", "ASSISTANT", "SYSTEM", name="chatrole"),
            autoincrement=False,
            nullable=False,
        ),
        sa.Column("tokens_used", sa.INTEGER(), autoincrement=False, nullable=True),
        sa.Column("organization_id", sa.UUID(), autoincrement=False, nullable=False),
        sa.Column("id", sa.UUID(), autoincrement=False, nullable=False),
        sa.Column("created_at", postgresql.TIMESTAMP(), autoincrement=False, nullable=False),
        sa.Column("modified_at", postgresql.TIMESTAMP(), autoincrement=False, nullable=False),
        sa.Column("created_by_email", sa.VARCHAR(), autoincrement=False, nullable=False),
        sa.Column("modified_by_email", sa.VARCHAR(), autoincrement=False, nullable=False),
        sa.ForeignKeyConstraint(
            ["chat_id"], ["chat.id"], name=op.f("chat_message_chat_id_fkey"), ondelete="CASCADE"
        ),
        sa.ForeignKeyConstraint(
            ["created_by_email"], ["user.email"], name=op.f("chat_message_created_by_email_fkey")
        ),
        sa.ForeignKeyConstraint(
            ["modified_by_email"], ["user.email"], name=op.f("chat_message_modified_by_email_fkey")
        ),
        sa.ForeignKeyConstraint(
            ["organization_id"], ["organization.id"], name=op.f("chat_message_organization_id_fkey")
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("chat_message_pkey")),
    )
    op.create_index(op.f("ix_chat_message_chat_id"), "chat_message", ["chat_id"], unique=False)
    op.create_table(
        "chat",
        sa.Column("name", sa.VARCHAR(length=255), autoincrement=False, nullable=False),
        sa.Column("sync_id", sa.UUID(), autoincrement=False, nullable=True),
        sa.Column("description", sa.TEXT(), autoincrement=False, nullable=True),
        sa.Column("model_name", sa.VARCHAR(length=50), autoincrement=False, nullable=False),
        sa.Column(
            "model_settings",
            postgresql.JSON(astext_type=sa.Text()),
            autoincrement=False,
            nullable=False,
        ),
        sa.Column(
            "search_settings",
            postgresql.JSON(astext_type=sa.Text()),
            autoincrement=False,
            nullable=False,
        ),
        sa.Column("organization_id", sa.UUID(), autoincrement=False, nullable=False),
        sa.Column("id", sa.UUID(), autoincrement=False, nullable=False),
        sa.Column("created_at", postgresql.TIMESTAMP(), autoincrement=False, nullable=False),
        sa.Column("modified_at", postgresql.TIMESTAMP(), autoincrement=False, nullable=False),
        sa.Column("created_by_email", sa.VARCHAR(), autoincrement=False, nullable=False),
        sa.Column("modified_by_email", sa.VARCHAR(), autoincrement=False, nullable=False),
        sa.ForeignKeyConstraint(
            ["created_by_email"], ["user.email"], name=op.f("chat_created_by_email_fkey")
        ),
        sa.ForeignKeyConstraint(
            ["modified_by_email"], ["user.email"], name=op.f("chat_modified_by_email_fkey")
        ),
        sa.ForeignKeyConstraint(
            ["organization_id"], ["organization.id"], name=op.f("chat_organization_id_fkey")
        ),
        sa.ForeignKeyConstraint(["sync_id"], ["sync.id"], name=op.f("chat_sync_id_fkey")),
        sa.PrimaryKeyConstraint("id", name=op.f("chat_pkey")),
    )
    op.drop_table("user_organization")
    # ### end Alembic commands ###
