"""Add owner_id column to Phase 2 tables for row-level ownership.

Adds nullable owner_id (FK → users.id, ON DELETE SET NULL) to tables
that gained OwnedMixin in Phase 2. For tables that already have a
non-nullable user_id column, owner_id is backfilled from user_id.

Revision ID: 0102
"""

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import UUID as PGUUID

revision = "0102"
down_revision = "0101"
branch_labels = None
depends_on = None

# (table_name, has_user_id_to_backfill)
TABLES = [
    # Core models
    ("contact_folders", True),
    ("user_preferences", True),
    ("workspaces", False),
    # Plugin models
    ("mcp_server_configs", False),
    ("automation_agent_definitions", False),
    ("automation_definitions", False),
    ("report_templates", False),
    ("report_instances", False),
    ("entity_links", False),
    ("comm_conversations", False),
    ("ai_proactive_suggestions", True),
    ("ai_agents", False),
    ("ai_chat_sessions", True),
    ("tags", False),
    ("share_links", True),
]


def upgrade() -> None:
    for table_name, _has_user_id in TABLES:
        op.add_column(
            table_name,
            sa.Column(
                "owner_id",
                PGUUID(as_uuid=True),
                sa.ForeignKey("users.id", ondelete="SET NULL"),
                nullable=True,
            ),
        )
        op.create_index(
            f"ix_{table_name}_owner_id",
            table_name,
            ["owner_id"],
        )

    # Backfill owner_id from user_id where available
    for table_name, has_user_id in TABLES:
        if has_user_id:
            op.execute(
                f"UPDATE {table_name} SET owner_id = user_id WHERE owner_id IS NULL;"
            )


def downgrade() -> None:
    for table_name, _ in TABLES:
        op.drop_index(f"ix_{table_name}_owner_id", table_name=table_name)
        op.drop_column(table_name, "owner_id")