"""add external qa followup state fields

Revision ID: 20260327_add_external_qa_followup_state
Revises: 20260327_cleanup_model_routing_legacy_names
Create Date: 2026-03-27
"""

from __future__ import annotations

from alembic import op
import sqlalchemy as sa

try:
    from sqlalchemy.dialects.postgresql import JSONB
except Exception:  # pragma: no cover
    JSONB = sa.JSON


revision = "20260327_add_external_qa_followup_state"
down_revision = "20260327_cleanup_model_routing_legacy_names"
branch_labels = None
depends_on = None


def _json_type(bind):
    return JSONB() if bind.dialect.name == "postgresql" else sa.JSON()


def _table_exists(bind, table_name: str) -> bool:
    return sa.inspect(bind).has_table(table_name)


def _column_exists(bind, table_name: str, column_name: str) -> bool:
    inspector = sa.inspect(bind)
    try:
        columns = inspector.get_columns(table_name)
    except Exception:
        return False
    return any(str(col.get("name")) == column_name for col in columns)


def _index_exists(bind, table_name: str, index_name: str) -> bool:
    inspector = sa.inspect(bind)
    try:
        indexes = inspector.get_indexes(table_name)
    except Exception:
        return False
    return any(idx.get("name") == index_name for idx in indexes)


def _create_index_if_not_exists(
    bind,
    index_name: str,
    table_name: str,
    columns: list[str],
    *,
    unique: bool = False,
) -> None:
    if not _index_exists(bind, table_name, index_name):
        op.create_index(index_name, table_name, columns, unique=unique)


def _drop_index_if_exists(bind, index_name: str, table_name: str) -> None:
    if _index_exists(bind, table_name, index_name):
        op.drop_index(index_name, table_name=table_name)


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

    if _table_exists(bind, "qa_sessions"):
        if not _column_exists(bind, "qa_sessions", "external_conversation_key"):
            with op.batch_alter_table("qa_sessions") as batch_op:
                batch_op.add_column(sa.Column("external_conversation_key", sa.VARCHAR(length=160), nullable=True))
        _create_index_if_not_exists(
            bind,
            "ix_qa_sessions_external_conversation_key",
            "qa_sessions",
            ["external_conversation_key"],
        )
        _create_index_if_not_exists(
            bind,
            "uq_qa_sessions_tenant_external_conversation_key",
            "qa_sessions",
            ["tenant_id", "external_conversation_key"],
            unique=True,
        )

    if _table_exists(bind, "qa_session_states"):
        if not _column_exists(bind, "qa_session_states", "pending_clarification"):
            with op.batch_alter_table("qa_session_states") as batch_op:
                batch_op.add_column(sa.Column("pending_clarification", _json_type(bind), nullable=True))
        if not _column_exists(bind, "qa_session_states", "pending_clarification_expires_at"):
            with op.batch_alter_table("qa_session_states") as batch_op:
                batch_op.add_column(sa.Column("pending_clarification_expires_at", sa.TIMESTAMP(), nullable=True))
        _create_index_if_not_exists(
            bind,
            "ix_qa_session_states_pending_clarification_expires_at",
            "qa_session_states",
            ["pending_clarification_expires_at"],
        )


def downgrade() -> None:
    bind = op.get_bind()

    if _table_exists(bind, "qa_session_states"):
        _drop_index_if_exists(
            bind,
            "ix_qa_session_states_pending_clarification_expires_at",
            "qa_session_states",
        )
        if _column_exists(bind, "qa_session_states", "pending_clarification_expires_at") or _column_exists(
            bind, "qa_session_states", "pending_clarification"
        ):
            with op.batch_alter_table("qa_session_states") as batch_op:
                if _column_exists(bind, "qa_session_states", "pending_clarification_expires_at"):
                    batch_op.drop_column("pending_clarification_expires_at")
                if _column_exists(bind, "qa_session_states", "pending_clarification"):
                    batch_op.drop_column("pending_clarification")

    if _table_exists(bind, "qa_sessions"):
        _drop_index_if_exists(
            bind,
            "uq_qa_sessions_tenant_external_conversation_key",
            "qa_sessions",
        )
        _drop_index_if_exists(
            bind,
            "ix_qa_sessions_external_conversation_key",
            "qa_sessions",
        )
        if _column_exists(bind, "qa_sessions", "external_conversation_key"):
            with op.batch_alter_table("qa_sessions") as batch_op:
                batch_op.drop_column("external_conversation_key")