"""resource scoped entitlements

Revision ID: 000000000005
Revises: 000000000004
Create Date: 2026-02-11

Adds resource_type and resource_id columns to the entitlements table
to support company, org, and system-scoped entitlements in addition
to user-scoped ones.
"""

from alembic import op

revision = "000000000005"
down_revision = "000000000004"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # Add resource_type column with default 'user'
    op.execute("""
        ALTER TABLE entitlements
        ADD COLUMN IF NOT EXISTS resource_type VARCHAR NOT NULL DEFAULT 'user';
    """)

    # Add CHECK constraint for allowed resource_type values
    op.execute("""
        ALTER TABLE entitlements
        ADD CONSTRAINT ck_entitlements_resource_type
        CHECK (resource_type IN ('user', 'company', 'system', 'org'));
    """)

    # Add resource_id column (nullable UUID)
    op.execute("""
        ALTER TABLE entitlements
        ADD COLUMN IF NOT EXISTS resource_id UUID;
    """)

    # Add CHECK constraint: resource_id required for company and org types
    op.execute("""
        ALTER TABLE entitlements
        ADD CONSTRAINT ck_entitlements_resource_id_required
        CHECK (
            resource_type NOT IN ('company', 'org')
            OR resource_id IS NOT NULL
        );
    """)

    # Add partial index for resource lookups
    op.execute("""
        CREATE INDEX IF NOT EXISTS idx_entitlements_resource
        ON entitlements (resource_type, resource_id)
        WHERE resource_id IS NOT NULL;
    """)

    # Add partial unique index for non-user resource scoping
    op.execute("""
        CREATE UNIQUE INDEX IF NOT EXISTS idx_entitlements_resource_unique
        ON entitlements (application_id, entitlement_key, resource_type, resource_id, source)
        WHERE resource_type != 'user';
    """)


def downgrade() -> None:
    op.execute("DROP INDEX IF EXISTS idx_entitlements_resource_unique;")
    op.execute("DROP INDEX IF EXISTS idx_entitlements_resource;")
    op.execute("""
        ALTER TABLE entitlements
        DROP CONSTRAINT IF EXISTS ck_entitlements_resource_id_required;
    """)
    op.execute("""
        ALTER TABLE entitlements
        DROP CONSTRAINT IF EXISTS ck_entitlements_resource_type;
    """)
    op.execute("ALTER TABLE entitlements DROP COLUMN IF EXISTS resource_id;")
    op.execute("ALTER TABLE entitlements DROP COLUMN IF EXISTS resource_type;")