"""Add embedding column to contacts + timestamp columns to audit_log.

Fixes two issues found by API integration tests:
1. contacts.embedding (vector(768)) — ORM model was updated in Phase 5.3 but
   the plugin migration 0002_embeddings.sql was never run as an Alembic migration.
2. audit_log.created_at, updated_at, deleted_at — AuditLog model inherits TenantMixin
   which expects these columns, but they were never added to the DB table.

Revision ID: 0104
"""

from alembic import op
import sqlalchemy as sa

revision = "0104"
down_revision = "0103"


def upgrade() -> None:
    conn = op.get_bind()

    # 0. Ensure pgvector extension is installed
    op.execute("CREATE EXTENSION IF NOT EXISTS vector")

    # 1. Add embedding column to contacts (if not exists)
    result = conn.execute(sa.text(
        "SELECT column_name FROM information_schema.columns "
        "WHERE table_name = 'contacts' AND column_name = 'embedding'"
    ))
    if result.fetchone() is None:
        op.execute("ALTER TABLE contacts ADD COLUMN embedding vector(768)")
        op.execute(
            "CREATE INDEX IF NOT EXISTS ix_contacts_embedding "
            "ON contacts USING hnsw(embedding vector_cosine_ops)"
        )

    # 2. Add embedding columns to other tables (from plugin migration 0002)
    for table in ["mails", "files", "calendar_entries"]:
        result = conn.execute(sa.text(
            f"SELECT column_name FROM information_schema.columns "
            f"WHERE table_name = '{table}' AND column_name = 'embedding'"
        ))
        if result.fetchone() is None:
            # Check if table exists
            table_exists = conn.execute(sa.text(
                f"SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = '{table}')"
            )).scalar()
            if table_exists:
                op.execute(f"ALTER TABLE {table} ADD COLUMN embedding vector(768)")
                op.execute(
                    f"CREATE INDEX IF NOT EXISTS ix_{table}_embedding "
                    f"ON {table} USING hnsw(embedding vector_cosine_ops)"
                )

    # Tags use 384-dim embeddings
    result = conn.execute(sa.text(
        "SELECT column_name FROM information_schema.columns "
        "WHERE table_name = 'tags' AND column_name = 'embedding'"
    ))
    if result.fetchone() is None:
        table_exists = conn.execute(sa.text(
            "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'tags')"
        )).scalar()
        if table_exists:
            op.execute("ALTER TABLE tags ADD COLUMN embedding vector(384)")
            op.execute(
                "CREATE INDEX IF NOT EXISTS ix_tags_embedding "
                "ON tags USING hnsw(embedding vector_cosine_ops)"
            )

    # 3. Add timestamp columns to audit_log (if not exists)
    for col in ["created_at", "updated_at", "deleted_at"]:
        result = conn.execute(sa.text(
            f"SELECT column_name FROM information_schema.columns "
            f"WHERE table_name = 'audit_log' AND column_name = '{col}'"
        ))
        if result.fetchone() is None:
            op.execute(
                f"ALTER TABLE audit_log ADD COLUMN {col} "
                f"TIMESTAMPTZ DEFAULT NOW()"
            )

    # 4. Add timestamp columns to deletion_log (if not exists)
    for col in ["created_at", "updated_at", "deleted_at"]:
        result = conn.execute(sa.text(
            f"SELECT column_name FROM information_schema.columns "
            f"WHERE table_name = 'deletion_log' AND column_name = '{col}'"
        ))
        if result.fetchone() is None:
            op.execute(
                f"ALTER TABLE deletion_log ADD COLUMN {col} "
                f"TIMESTAMPTZ DEFAULT NOW()"
            )


def downgrade() -> None:
    # Drop embedding columns
    for table in ["contacts", "mails", "companies", "files", "calendar_entries"]:
        op.execute(f"DROP INDEX IF EXISTS ix_{table}_embedding")
        op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS embedding")

    op.execute("DROP INDEX IF EXISTS ix_tags_embedding")
    op.execute("ALTER TABLE tags DROP COLUMN IF EXISTS embedding")

    # Drop timestamp columns from audit_log
    for col in ["created_at", "updated_at", "deleted_at"]:
        op.execute(f"ALTER TABLE audit_log DROP COLUMN IF EXISTS {col}")

    # Drop timestamp columns from deletion_log
    for col in ["created_at", "updated_at", "deleted_at"]:
        op.execute(f"ALTER TABLE deletion_log DROP COLUMN IF EXISTS {col}")