"""extend prompt eval replay

Revision ID: 20260430_extend_prompt_eval_replay
Revises: 20260430_add_prompt_replay_runs
Create Date: 2026-04-30
"""

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 = "20260430_extend_prompt_eval_replay"
down_revision = "20260430_add_prompt_replay_runs"
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:
    if not _table_exists(bind, table_name):
        return False
    return column_name in {str(column.get("name") or "") for column in sa.inspect(bind).get_columns(table_name)}


def _index_exists(bind, table_name: str, index_name: str) -> bool:
    if not _table_exists(bind, table_name):
        return False
    return index_name in {str(index.get("name") or "") for index in sa.inspect(bind).get_indexes(table_name)}


def _add_column_if_missing(bind, table_name: str, column: sa.Column) -> None:
    if not _column_exists(bind, table_name, str(column.name)):
        op.add_column(table_name, column)


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


def upgrade() -> None:
    bind = op.get_bind()
    if _table_exists(bind, "prompt_golden_cases"):
        _add_column_if_missing(bind, "prompt_golden_cases", sa.Column("case_type", sa.String(length=32), nullable=False, server_default="regression"))
        _add_column_if_missing(bind, "prompt_golden_cases", sa.Column("priority", sa.String(length=24), nullable=False, server_default="blocking"))
        _add_column_if_missing(bind, "prompt_golden_cases", sa.Column("scenario_key", sa.String(length=128), nullable=True))
        _add_column_if_missing(bind, "prompt_golden_cases", sa.Column("tags", _json_type(bind), nullable=True))
        _add_column_if_missing(bind, "prompt_golden_cases", sa.Column("last_run_status", sa.String(length=24), nullable=True))
        _add_column_if_missing(bind, "prompt_golden_cases", sa.Column("last_failed_at", sa.String(length=64), nullable=True))
        _create_index_if_missing(bind, "ix_prompt_golden_cases_case_type", "prompt_golden_cases", ["case_type"])
        _create_index_if_missing(bind, "ix_prompt_golden_cases_priority", "prompt_golden_cases", ["priority"])
        _create_index_if_missing(bind, "ix_prompt_golden_cases_scenario_key", "prompt_golden_cases", ["scenario_key"])
        _create_index_if_missing(bind, "ix_prompt_golden_cases_last_run_status", "prompt_golden_cases", ["last_run_status"])

    if _table_exists(bind, "prompt_replay_case_results"):
        _add_column_if_missing(bind, "prompt_replay_case_results", sa.Column("baseline_ai_call_log_id", sa.Integer(), nullable=True))
        _add_column_if_missing(bind, "prompt_replay_case_results", sa.Column("candidate_ai_call_log_id", sa.Integer(), nullable=True))
        _add_column_if_missing(bind, "prompt_replay_case_results", sa.Column("judge_score", sa.Float(), nullable=True))
        _add_column_if_missing(bind, "prompt_replay_case_results", sa.Column("judge_label", sa.String(length=32), nullable=True))
        _add_column_if_missing(bind, "prompt_replay_case_results", sa.Column("judge_reason", sa.Text(), nullable=True))
        _create_index_if_missing(bind, "ix_prompt_replay_case_results_judge_label", "prompt_replay_case_results", ["judge_label"])
        _create_index_if_missing(bind, "ix_prompt_replay_case_results_baseline_ai_call_log_id", "prompt_replay_case_results", ["baseline_ai_call_log_id"])
        _create_index_if_missing(bind, "ix_prompt_replay_case_results_candidate_ai_call_log_id", "prompt_replay_case_results", ["candidate_ai_call_log_id"])


def downgrade() -> None:
    bind = op.get_bind()
    if _table_exists(bind, "prompt_replay_case_results"):
        for index_name in (
            "ix_prompt_replay_case_results_candidate_ai_call_log_id",
            "ix_prompt_replay_case_results_baseline_ai_call_log_id",
            "ix_prompt_replay_case_results_judge_label",
        ):
            if _index_exists(bind, "prompt_replay_case_results", index_name):
                op.drop_index(index_name, table_name="prompt_replay_case_results")
        for column_name in ("judge_reason", "judge_label", "judge_score", "candidate_ai_call_log_id", "baseline_ai_call_log_id"):
            if _column_exists(bind, "prompt_replay_case_results", column_name):
                op.drop_column("prompt_replay_case_results", column_name)
    if _table_exists(bind, "prompt_golden_cases"):
        for index_name in (
            "ix_prompt_golden_cases_last_run_status",
            "ix_prompt_golden_cases_scenario_key",
            "ix_prompt_golden_cases_priority",
            "ix_prompt_golden_cases_case_type",
        ):
            if _index_exists(bind, "prompt_golden_cases", index_name):
                op.drop_index(index_name, table_name="prompt_golden_cases")
        for column_name in ("last_failed_at", "last_run_status", "tags", "scenario_key", "priority", "case_type"):
            if _column_exists(bind, "prompt_golden_cases", column_name):
                op.drop_column("prompt_golden_cases", column_name)