"""migrate_ai_model_to_claude

Revision ID: 3da4f2d49e3e
Revises: 84deefdcf6c4
Create Date: 2025-04-23 10:21:07.733928

"""
import json
from alembic import op
import sqlalchemy as sa

from models.user_preference import Preference

# revision identifiers, used by Alembic.
revision = '3da4f2d49e3e'
down_revision = '84deefdcf6c4'
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    # loop all the preference, filter items which preference.preference ai_model contains "gpt-3.5-turbo", replace it with "claude-sonnet"
    conn = op.get_bind()
    preferences = conn.execute(sa.text("SELECT id, preference FROM preference"))
    for i in preferences:
        preference = json.loads(i.preference)
        llm_model = preference.get("llm_model", "") or preference.get("ai_model", "")
        mode_type = preference.get("mode_type", "")
        llm_api_key = preference.get("llm_api_key", "") or preference.get("openai_api_key", "")

        if mode_type:
            print(f"skip {i.id} because mode_type is not empty")
            continue

        if llm_api_key:
            print(f"skip {i.id} because llm_api_key is not empty")
            continue

        if llm_model != "gpt-3.5-turbo":
            print(f"skip {i.id} because llm_model is not gpt-3.5-turbo")
            continue

        preference["llm_model"] = "claude-sonnet"
        conn.execute(sa.text("UPDATE preference SET preference = :preference WHERE id = :id"), {"preference": json.dumps(preference), "id": i.id})
        print(f"update {i.id} to claude-sonnet")
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###