"""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 table exists
        table_exists = conn.execute(
            sa.text(
                "SELECT 1 FROM information_schema.tables WHERE table_name = :table"
            ),
            {"table": table},
        ).fetchone()
        if table_exists is None:
            continue
        # 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.execute(f"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS owner_id UUID REFERENCES users(id) ON DELETE SET NULL")
            op.execute(f"CREATE INDEX IF NOT EXISTS ix_{table}_owner ON {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