"""Add KnowledgePack, KnowledgeAgentRelationship and add kp_id to Conversation

Revision ID: b0e89352628b
Revises: 898918bedceb
Create Date: 2025-09-23 16:26:06.843442

"""

from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = "b0e89352628b"
down_revision: Union[str, Sequence[str], None] = "898918bedceb"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    """Upgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "knowledge_packs",
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column("metadata", sa.JSON(), nullable=True),
        sa.Column("created_at", sa.DateTime(), nullable=True),
        sa.Column("updated_at", sa.DateTime(), nullable=True),
        sa.Column("source_kp_id", sa.Integer(), nullable=True),
        sa.ForeignKeyConstraint(["source_kp_id"], ["knowledge_packs.id"], name="fk_knowledge_packs_source_kp_id"),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(op.f("ix_knowledge_packs_id"), "knowledge_packs", ["id"], unique=False)
    op.create_index(op.f("ix_knowledge_packs_source_kp_id"), "knowledge_packs", ["source_kp_id"], unique=False)
    # Note: source_kp_id creates a self-referential relationship for knowledge pack hierarchy
    op.create_table(
        "conversations_v2",
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column("title", sa.String(), nullable=False),
        sa.Column("agent_id", sa.Integer(), nullable=True),
        sa.Column("kp_id", sa.Integer(), nullable=True),
        sa.Column("code_gen_id", sa.Integer(), nullable=True),
        sa.Column("created_at", sa.DateTime(), nullable=True),
        sa.Column("updated_at", sa.DateTime(), nullable=True),
        sa.ForeignKeyConstraint(
            ["agent_id"],
            ["agents.id"],
        ),
        sa.ForeignKeyConstraint(
            ["code_gen_id"],
            ["agents.id"],
        ),
        sa.ForeignKeyConstraint(
            ["kp_id"],
            ["knowledge_packs.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(op.f("ix_conversations_v2_agent_id"), "conversations_v2", ["agent_id"], unique=False)
    op.create_index(op.f("ix_conversations_v2_code_gen_id"), "conversations_v2", ["code_gen_id"], unique=False)
    op.create_index(op.f("ix_conversations_v2_id"), "conversations_v2", ["id"], unique=False)
    op.create_index(op.f("ix_conversations_v2_kp_id"), "conversations_v2", ["kp_id"], unique=False)
    op.create_table(
        "knowledge_agent_relationships",
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column("knowledge_pack_id", sa.Integer(), nullable=False),
        sa.Column("agent_id", sa.Integer(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=True),
        sa.Column("updated_at", sa.DateTime(), nullable=True),
        sa.ForeignKeyConstraint(
            ["agent_id"],
            ["agents.id"],
        ),
        sa.ForeignKeyConstraint(
            ["knowledge_pack_id"],
            ["knowledge_packs.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(op.f("ix_knowledge_agent_relationships_agent_id"), "knowledge_agent_relationships", ["agent_id"], unique=False)
    op.create_index(op.f("ix_knowledge_agent_relationships_id"), "knowledge_agent_relationships", ["id"], unique=False)
    op.create_index(
        op.f("ix_knowledge_agent_relationships_knowledge_pack_id"), "knowledge_agent_relationships", ["knowledge_pack_id"], unique=False
    )
    op.create_table(
        "messages_v2",
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column("conversation_id", sa.Integer(), nullable=False),
        sa.Column("sender", sa.String(), nullable=False),
        sa.Column("content", sa.Text(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=True),
        sa.Column("updated_at", sa.DateTime(), nullable=True),
        sa.ForeignKeyConstraint(
            ["conversation_id"],
            ["conversations_v2.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(op.f("ix_messages_v2_conversation_id"), "messages_v2", ["conversation_id"], unique=False)
    op.create_index(op.f("ix_messages_v2_id"), "messages_v2", ["id"], unique=False)

    # ### Data migration from old tables to new v2 tables ###
    # Migrate conversations to conversations_v2
    # Note: kp_id and code_gen_id are set to NULL since old conversations don't have these fields
    op.execute("""
        INSERT INTO conversations_v2 (id, title, agent_id, kp_id, code_gen_id, created_at, updated_at)
        SELECT id, title, agent_id, NULL as kp_id, NULL as code_gen_id, created_at, updated_at
        FROM conversations
    """)

    # Migrate messages to messages_v2
    # Note: conversation_id references will be maintained since we preserve the same IDs
    op.execute("""
        INSERT INTO messages_v2 (id, conversation_id, sender, content, created_at, updated_at)
        SELECT id, conversation_id, sender, content, created_at, updated_at
        FROM messages
    """)

    # Drop old tables and their indexes
    op.drop_index(op.f("ix_messages_id"), table_name="messages")
    op.drop_index(op.f("ix_messages_conversation_id"), table_name="messages")
    op.drop_table("messages")
    op.drop_index(op.f("ix_conversations_id"), table_name="conversations")
    op.drop_index(op.f("ix_conversations_agent_id"), table_name="conversations")
    op.drop_table("conversations")
    # ### end Alembic commands ###


def downgrade() -> None:
    """Downgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###

    # ### Recreate old tables for downgrade ###
    op.create_table(
        "conversations",
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column("title", sa.String(), nullable=False),
        sa.Column("agent_id", sa.Integer(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=True),
        sa.Column("updated_at", sa.DateTime(), nullable=True),
        sa.ForeignKeyConstraint(
            ["agent_id"],
            ["agents.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(op.f("ix_conversations_agent_id"), "conversations", ["agent_id"], unique=False)
    op.create_index(op.f("ix_conversations_id"), "conversations", ["id"], unique=False)

    op.create_table(
        "messages",
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column("conversation_id", sa.Integer(), nullable=False),
        sa.Column("sender", sa.String(), nullable=False),
        sa.Column("content", sa.Text(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=True),
        sa.Column("updated_at", sa.DateTime(), nullable=True),
        sa.ForeignKeyConstraint(
            ["conversation_id"],
            ["conversations.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(op.f("ix_messages_conversation_id"), "messages", ["conversation_id"], unique=False)
    op.create_index(op.f("ix_messages_id"), "messages", ["id"], unique=False)

    # ### Migrate data back from v2 tables to old tables ###
    # Migrate conversations_v2 back to conversations
    # Note: Only migrate conversations that have agent_id (ignore knowledge pack only conversations)
    # Note: code_gen_id is ignored since old conversations table doesn't support it
    op.execute("""
        INSERT INTO conversations (id, title, agent_id, created_at, updated_at)
        SELECT id, title, agent_id, created_at, updated_at
        FROM conversations_v2
        WHERE agent_id IS NOT NULL
    """)

    # Migrate messages_v2 back to messages
    # Note: conversation_id references will be maintained since we preserve the same IDs
    op.execute("""
        INSERT INTO messages (id, conversation_id, sender, content, created_at, updated_at)
        SELECT id, conversation_id, sender, content, created_at, updated_at
        FROM messages_v2
    """)

    # Drop v2 tables and their indexes
    op.drop_index(op.f("ix_messages_v2_id"), table_name="messages_v2")
    op.drop_index(op.f("ix_messages_v2_conversation_id"), table_name="messages_v2")
    op.drop_table("messages_v2")
    op.drop_index(op.f("ix_knowledge_agent_relationships_knowledge_pack_id"), table_name="knowledge_agent_relationships")
    op.drop_index(op.f("ix_knowledge_agent_relationships_id"), table_name="knowledge_agent_relationships")
    op.drop_index(op.f("ix_knowledge_agent_relationships_agent_id"), table_name="knowledge_agent_relationships")
    op.drop_table("knowledge_agent_relationships")
    op.drop_index(op.f("ix_conversations_v2_kp_id"), table_name="conversations_v2")
    op.drop_index(op.f("ix_conversations_v2_id"), table_name="conversations_v2")
    op.drop_index(op.f("ix_conversations_v2_code_gen_id"), table_name="conversations_v2")
    op.drop_index(op.f("ix_conversations_v2_agent_id"), table_name="conversations_v2")
    op.drop_table("conversations_v2")
    op.drop_index(op.f("ix_knowledge_packs_source_kp_id"), table_name="knowledge_packs")
    op.drop_index(op.f("ix_knowledge_packs_id"), table_name="knowledge_packs")
    op.drop_table("knowledge_packs")
    # ### end Alembic commands ###
