"""cleanup legacy model routing function names

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

from __future__ import annotations

from copy import deepcopy
from typing import Any

from alembic import op
import sqlalchemy as sa


revision = "20260327_cleanup_model_routing_legacy_names"
down_revision = "20260326_add_ai_call_log_routing_fields"
branch_labels = None
depends_on = None

_RENAME_KEYS = {
    "project_draft_analysis_glm47_fallback": "ai_pmo_signal_title_refine",
    "adaptive_rag_answer_deep_signal_batch_analysis": "ai_pmo_signal_batch_analysis",
}
_REMOVE_KEYS = {
    "project_draft_charter_glm47_fallback",
    "project_draft_signal_glm47_fallback",
    "project_draft_task_chain_glm47_fallback",
}


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


def _integration_table(bind) -> sa.Table:
    metadata = sa.MetaData()
    return sa.Table(
        "integrations",
        metadata,
        sa.Column("id", sa.Integer(), primary_key=True),
        sa.Column("type", sa.String(length=100)),
        sa.Column("config", sa.JSON()),
    )


def _safe_dict(value: Any) -> dict[str, Any]:
    if isinstance(value, dict):
        try:
            return deepcopy(value)
        except Exception:
            return dict(value)
    return {}


def _rewrite_model_routing_config(config: Any) -> dict[str, Any]:
    payload = _safe_dict(config)
    if not payload:
        return {}

    mappings = payload.get("mappings")
    if not isinstance(mappings, dict):
        payload["mappings"] = {}
        return payload

    normalized: dict[str, Any] = {}
    for key, value in mappings.items():
        name = str(key or "").strip().lower()
        if not name:
            continue
        normalized[name] = deepcopy(value)

    for old_key, new_key in _RENAME_KEYS.items():
        if old_key not in normalized:
            continue
        if new_key not in normalized:
            normalized[new_key] = normalized[old_key]
        normalized.pop(old_key, None)

    for key in _REMOVE_KEYS:
        normalized.pop(key, None)

    payload["mappings"] = normalized
    return payload


def upgrade() -> None:
    bind = op.get_bind()
    if not _table_exists(bind, "integrations"):
        return

    table = _integration_table(bind)
    rows = (
        bind.execute(
            sa.select(table.c.id, table.c.config).where(table.c.type == "model_routing")
        )
        .mappings()
        .all()
    )
    for row in rows:
        row_id = int(row.get("id") or 0)
        rewritten = _rewrite_model_routing_config(row.get("config"))
        if not rewritten:
            continue
        if rewritten == row.get("config"):
            continue
        bind.execute(
            table.update().where(table.c.id == row_id).values(config=rewritten)
        )


def downgrade() -> None:
    bind = op.get_bind()
    if not _table_exists(bind, "integrations"):
        return

    table = _integration_table(bind)
    reverse_rename = {new_key: old_key for old_key, new_key in _RENAME_KEYS.items()}
    rows = (
        bind.execute(
            sa.select(table.c.id, table.c.config).where(table.c.type == "model_routing")
        )
        .mappings()
        .all()
    )
    for row in rows:
        row_id = int(row.get("id") or 0)
        payload = _safe_dict(row.get("config"))
        mappings = payload.get("mappings")
        if not isinstance(mappings, dict):
            continue
        normalized: dict[str, Any] = {}
        for key, value in mappings.items():
            name = str(key or "").strip().lower()
            if not name:
                continue
            normalized[name] = deepcopy(value)
        for new_key, old_key in reverse_rename.items():
            if new_key not in normalized:
                continue
            if old_key not in normalized:
                normalized[old_key] = normalized[new_key]
            normalized.pop(new_key, None)
        payload["mappings"] = normalized
        if payload == row.get("config"):
            continue
        bind.execute(
            table.update().where(table.c.id == row_id).values(config=payload)
        )