from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import JSONB, UUID

revision = "026"
down_revision = "025"
branch_labels = None
depends_on = None


def upgrade() -> None:
    bind = op.get_bind()
    ctx = op.get_context()

    # PostgreSQL requires newly added enum values to be committed
    # before they are referenced in DML.
    with ctx.autocommit_block():
        bind.execute(
            sa.text(
                "ALTER TYPE p2hs_codegen_run_status "
                "ADD VALUE IF NOT EXISTS 'clarification_required'"
            )
        )
        bind.execute(
            sa.text(
                "ALTER TYPE p2hs_codegen_run_status "
                "ADD VALUE IF NOT EXISTS 'user_input_required'"
            )
        )

    op.execute(
        sa.text(
            "UPDATE codegen_runs "
            "SET status = 'user_input_required' "
            "WHERE status = 'awaiting_user_input'"
        )
    )

    op.create_table(
        "run_chat_messages",
        sa.Column(
            "id",
            UUID(as_uuid=True),
            primary_key=True,
            server_default=sa.text("gen_random_uuid()"),
        ),
        sa.Column(
            "project_id",
            UUID(as_uuid=True),
            sa.ForeignKey("projects.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column(
            "session_id",
            UUID(as_uuid=True),
            sa.ForeignKey("prompt_sessions.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column(
            "run_id",
            UUID(as_uuid=True),
            sa.ForeignKey("codegen_runs.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column("role", sa.String(16), nullable=False),
        sa.Column("message_type", sa.String(32), nullable=False, server_default="message"),
        sa.Column("content", sa.Text(), nullable=False),
        sa.Column("payload", JSONB, nullable=True),
        sa.Column(
            "prompt_id",
            UUID(as_uuid=True),
            sa.ForeignKey("prompts.id", ondelete="SET NULL"),
            nullable=True,
        ),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
    )
    op.create_index("ix_run_chat_messages_session_id", "run_chat_messages", ["session_id"])
    op.create_index("ix_run_chat_messages_run_id", "run_chat_messages", ["run_id"])
    op.create_index("ix_run_chat_messages_project_id", "run_chat_messages", ["project_id"])