"""Add relation_type and metadata columns to container_relations

Revision ID: relations_0002
Revises: relations_0001
Create Date: 2026-02-28
"""

import sqlalchemy as sa
from alembic import op

revision = "relations_0002"
down_revision = "relations_0001"
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.add_column(
        "container_relations",
        sa.Column("relation_type", sa.String(64), nullable=True),
    )
    op.add_column(
        "container_relations",
        sa.Column("metadata", sa.JSON(), nullable=True),
    )
    op.create_index(
        "ix_container_relations_relation_type",
        "container_relations",
        ["relation_type", "parent_type", "parent_id"],
    )

    # Backfill relation_type for existing rows based on parent/child type pairs
    op.execute("""
        UPDATE container_relations
        SET relation_type = CASE
            WHEN parent_type = 'page' AND child_type = 'market'
                THEN 'page->market'
            WHEN parent_type = 'page' AND child_type = 'calendar'
                THEN 'page->calendar'
            WHEN parent_type = 'page' AND child_type = 'menu_node'
                THEN 'page->menu_node'
        END
        WHERE relation_type IS NULL
          AND parent_type = 'page'
          AND child_type IN ('market', 'calendar', 'menu_node')
    """)


def downgrade() -> None:
    op.drop_index("ix_container_relations_relation_type", "container_relations")
    op.drop_column("container_relations", "metadata")
    op.drop_column("container_relations", "relation_type")
