"""Add owner_id to plugin entity tables for row-level ownership.

Revision ID: 0054
Revises: 0053
Create Date: 2026-07-29
"""

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

revision = "0054"
down_revision = "0053"
branch_labels = None
depends_on = None

# Tables that need owner_id
TABLES = [
    "files",
    "folders",
    "calendar_entries",
    "calendars",
    "tasks",
    "subtasks",
]


def upgrade() -> None:
    # Check which columns already exist before adding
    conn = op.get_bind()
    for table in TABLES:
        # Check if column already exists
        result = conn.execute(
            sa.text(
                "SELECT column_name FROM information_schema.columns "
                "WHERE table_name = :table AND column_name = 'owner_id'"
            ),
            {"table": table},
        )
        if result.fetchone() is None:
            op.add_column(
                table,
                sa.Column(
                    "owner_id",
                    PGUUID(as_uuid=True),
                    sa.ForeignKey("users.id", ondelete="SET NULL"),
                    nullable=True,
                ),
            )
            op.create_index(f"ix_{table}_owner", table, ["owner_id"])


def downgrade() -> None:
    for table in TABLES:
        try:
            op.drop_index(f"ix_{table}_owner", table_name=table)
        except Exception:
            pass
        try:
            op.drop_column(table, "owner_id")
        except Exception:
            pass