"""extend tags for vision pipeline

Revision ID: 0002_extend_tags
Revises: 0001_baseline
Create Date: 2026-04-10

Add kind, source, representative_photo_id to tags table.
Add confidence, bbox, source to photo_tags association.
Switch uniqueness from (name) to (name, kind).
"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import JSONB

revision: str = "0002_extend_tags"
down_revision: Union[str, None] = "0001_baseline"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ── tags table ────────────────────────────────────────────────────
    op.execute("ALTER TABLE tags ADD COLUMN IF NOT EXISTS kind VARCHAR NOT NULL DEFAULT 'user'")
    op.execute("ALTER TABLE tags ADD COLUMN IF NOT EXISTS source VARCHAR")
    op.execute("ALTER TABLE tags ADD COLUMN IF NOT EXISTS representative_photo_id VARCHAR REFERENCES photos(id) ON DELETE SET NULL")

    # Create index on kind for filtering
    op.execute("CREATE INDEX IF NOT EXISTS ix_tags_kind ON tags(kind)")

    # Drop old unique constraint on name (if it exists) and add (name, kind).
    # SQLAlchemy create_all may have created either — handle both cases.
    op.execute("""
        DO $$
        BEGIN
            -- Drop the old single-column unique index/constraint if present.
            IF EXISTS (
                SELECT 1 FROM pg_indexes
                WHERE tablename = 'tags' AND indexname = 'ix_tags_name'
            ) THEN
                DROP INDEX ix_tags_name;
            END IF;

            -- Some SQLAlchemy versions create a unique constraint directly.
            IF EXISTS (
                SELECT 1 FROM information_schema.table_constraints
                WHERE table_name = 'tags' AND constraint_name = 'tags_name_key'
            ) THEN
                ALTER TABLE tags DROP CONSTRAINT tags_name_key;
            END IF;
        END $$;
    """)

    op.execute("""
        DO $$
        BEGIN
            IF NOT EXISTS (
                SELECT 1 FROM pg_constraint WHERE conname = 'uq_tags_name_kind'
            ) THEN
                ALTER TABLE tags ADD CONSTRAINT uq_tags_name_kind UNIQUE (name, kind);
            END IF;
        END $$;
    """)

    # ── photo_tags table ──────────────────────────────────────────────
    op.execute("ALTER TABLE photo_tags ADD COLUMN IF NOT EXISTS confidence FLOAT")
    op.execute("ALTER TABLE photo_tags ADD COLUMN IF NOT EXISTS bbox JSONB")
    op.execute("ALTER TABLE photo_tags ADD COLUMN IF NOT EXISTS source VARCHAR")


def downgrade() -> None:
    # photo_tags columns
    op.drop_column("photo_tags", "source")
    op.drop_column("photo_tags", "bbox")
    op.drop_column("photo_tags", "confidence")

    # tags: restore old unique constraint
    op.execute("ALTER TABLE tags DROP CONSTRAINT IF EXISTS uq_tags_name_kind")
    op.execute("CREATE UNIQUE INDEX IF NOT EXISTS ix_tags_name ON tags(name)")

    # tags columns
    op.drop_column("tags", "representative_photo_id")
    op.drop_column("tags", "source")
    op.drop_column("tags", "kind")
