"""Add require_user, treat_as_tool, metadata to messages_v2

Revision ID: 3364c5b025b4
Revises: b0e89352628b
Create Date: 2025-09-23 23:02:23.179688

"""

from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = "3364c5b025b4"
down_revision: Union[str, Sequence[str], None] = "b0e89352628b"
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! ###

    # SQLite doesn't support adding NOT NULL columns with default values to existing tables
    # So we need to recreate the table with the new columns and proper constraints

    # Create new table with NOT NULL constraints and default values
    op.create_table(
        "messages_v2_new",
        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("require_user", sa.Boolean(), nullable=False, default=False),
        sa.Column("treat_as_tool", sa.Boolean(), nullable=False, default=False),
        sa.Column("metadata", sa.JSON(), nullable=False, default={}),
        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"),
    )

    # Copy data from old table to new table with default values for new columns
    op.execute("""
        INSERT INTO messages_v2_new (id, conversation_id, sender, content, require_user, treat_as_tool, metadata, created_at, updated_at)
        SELECT id, conversation_id, sender, content, 
               false as require_user,
               false as treat_as_tool,
               '{}' as metadata,
               created_at, updated_at
        FROM messages_v2
    """)

    # Drop old table and rename new table
    op.drop_table("messages_v2")
    op.rename_table("messages_v2_new", "messages_v2")

    # Recreate indexes
    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)
    # ### end Alembic commands ###


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

    # Recreate the original messages_v2 table without the new columns
    op.create_table(
        "messages_v2_old",
        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"),
    )

    # Copy data from current table to old table (excluding new columns)
    op.execute("""
        INSERT INTO messages_v2_old (id, conversation_id, sender, content, created_at, updated_at)
        SELECT id, conversation_id, sender, content, created_at, updated_at
        FROM messages_v2
    """)

    # Drop current table and rename old table
    op.drop_table("messages_v2")
    op.rename_table("messages_v2_old", "messages_v2")

    # Recreate indexes
    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)
    # ### end Alembic commands ###
