"""Add is_multimodal column to llm_providers with heuristic backfill

Revision ID: 9f8e7d6c5b4a
Revises: b1c2d3e4f5a6
Create Date: 2026-05-20 14:30:00.000000

"""

from collections.abc import Sequence

import sqlalchemy as sa

from alembic import op
from sqlalchemy import inspect

# revision identifiers, used by Alembic.
revision: str = "9f8e7d6c5b4a"
down_revision: str | Sequence[str] | None = "b1c2d3e4f5a6"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def _is_multimodal_heuristic(model_name: str) -> bool:
    """Backfill heuristic: mirrors the logic from services/llm_prompt_templates.py
    at the time of migration creation. Returns True for known multimodal/vision models.
    """
    token = (model_name or "").strip().lower()
    if not token:
        return False
    text_only_markers = (
        "text-only",
        "embedding",
        "rerank",
        "whisper",
        "tts",
        "transcribe",
    )
    if any(marker in token for marker in text_only_markers):
        return False
    multimodal_markers = (
        "gpt-4o",
        "gpt-4.1",
        "vision",
        "vl",
        "gemini",
        "claude-3",
        "claude-sonnet-4",
        "omni",
        "multimodal",
        "pixtral",
        "llava",
        "internvl",
        "qwen-vl",
        "qwen2-vl",
    )
    return any(marker in token for marker in multimodal_markers)


def _table_exists(table_name: str) -> bool:
    """Check if a table exists in the current database schema."""
    conn = op.get_bind()
    inspector = inspect(conn)
    return table_name in inspector.get_table_names()


def upgrade() -> None:
    """Add is_multimodal column and backfill using heuristic."""
    # The llm_providers table is defined in the ORM model (src/models/llm.py)
    # and created by Base.metadata.create_all() at app startup, not via migrations.
    # If the table doesn't exist yet (fresh DB), skip — the model definition
    # already includes is_multimodal, so create_all() will create it with the column.
    if not _table_exists("llm_providers"):
        return

    # Step 1: Add column as nullable first
    op.add_column(
        "llm_providers",
        sa.Column("is_multimodal", sa.Boolean(), nullable=True),
    )

    # Step 2: Backfill existing rows using heuristic on default_model
    connection = op.get_bind()
    providers = connection.execute(
        sa.text("SELECT id, default_model FROM llm_providers")
    ).fetchall()

    for row in providers:
        multimodal = _is_multimodal_heuristic(row.default_model)
        connection.execute(
            sa.text(
                "UPDATE llm_providers SET is_multimodal = :multimodal WHERE id = :id"
            ),
            {"multimodal": multimodal, "id": row.id},
        )

    # Step 3: Make column non-nullable with default False
    # Use text("false") for PostgreSQL compatibility (not "0")
    op.alter_column("llm_providers", "is_multimodal", nullable=False, server_default=sa.text("false"))


def downgrade() -> None:
    """Remove is_multimodal column."""
    op.drop_column("llm_providers", "is_multimodal")
