"""Add plan mode, visual edit, and websocket support to codegen_runs.

Adds:
  - plan_ready status to the codegen run status enum
  - run_mode column (agent | plan | visual_edit)
  - plan_document column (markdown plan content for plan mode)
  - parent_plan_id column (FK to self — links agent execution back to approved plan)
  - selected_elements column (JSONB element metadata for visual edit mode)

Revision ID: 020
Revises: 019
Create Date: 2026-03-31
"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import JSONB, UUID as PG_UUID

revision = "020"
down_revision = "019"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # PostgreSQL requires ALTER TYPE outside a transaction block.
    # 'plan_ready' is the terminal status for plan mode runs (before user approves).
    op.execute("ALTER TYPE p2hs_codegen_run_status ADD VALUE IF NOT EXISTS 'plan_ready'")

    op.add_column(
        "codegen_runs",
        sa.Column("run_mode", sa.String(20), nullable=False, server_default="agent"),
    )
    op.add_column(
        "codegen_runs",
        sa.Column("plan_document", sa.Text(), nullable=True),
    )
    op.add_column(
        "codegen_runs",
        sa.Column(
            "parent_plan_id",
            PG_UUID(as_uuid=True),
            nullable=True,
        ),
    )
    op.create_foreign_key(
        "fk_codegen_runs_parent_plan_id",
        "codegen_runs",
        "codegen_runs",
        ["parent_plan_id"],
        ["id"],
        ondelete="SET NULL",
    )
    op.add_column(
        "codegen_runs",
        sa.Column("selected_elements", JSONB, nullable=True),
    )


def downgrade() -> None:
    op.drop_constraint("fk_codegen_runs_parent_plan_id", "codegen_runs", type_="foreignkey")
    op.drop_column("codegen_runs", "selected_elements")
    op.drop_column("codegen_runs", "parent_plan_id")
    op.drop_column("codegen_runs", "plan_document")
    op.drop_column("codegen_runs", "run_mode")
    # PostgreSQL does not support removing enum values — downgrade leaves plan_ready in the type.