"""refactor_agent_model config

Revision ID: 69cb68688942
Revises: 031b34b4dd30
Create Date: 2026-01-31 16:52:49.151030

"""
import json
import logging
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import sqlite

from openjiuwen_studio.core.database.migration_utils import table_exists, column_exists, index_exists

# revision identifiers, used by Alembic.
revision: str = '69cb68688942'
down_revision: Union[str, Sequence[str], None] = '031b34b4dd30'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    """Upgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    # 1. 首先添加新列
    for table in ['agent', 'agent_publish']:
        with op.batch_alter_table(table, schema=None) as batch_op:
            if not column_exists(table, 'model_id'):
                batch_op.add_column(sa.Column('model_id', sa.Integer(), nullable=True))
            if not column_exists(table, 'agent_model_config'):
                batch_op.add_column(sa.Column('agent_model_config', sa.JSON(), nullable=True))

    # 2. 数据迁移
    conn = op.get_bind()
    for table in ['agent', 'agent_publish']:
        if column_exists(table, 'model'):
            try:
                res = conn.execute(sa.text(f"SELECT id, model FROM {table} WHERE model IS NOT NULL"))
                for row in res:
                    row_id = row[0]
                    model_data_str = row[1]
                    if not model_data_str:
                        continue
                    
                    try:
                        if isinstance(model_data_str, str):
                            model_data = json.loads(model_data_str)
                        else:
                            model_data = model_data_str
                        
                        model_info = model_data.get("model_info", {})
                        model_id_val = model_info.get("model_id")
                        
                        if model_id_val is not None:
                            # 提取特定的配置字段，保持与 AgentModelConfig 结构一致
                            config_to_save = {}
                            for field in ["temperature", "top_p", "timeout"]:
                                if field in model_info:
                                    config_to_save[field] = model_info[field]
                                    
                            conn.execute(
                                sa.text(
                                    f"UPDATE {table} SET model_id = :mid, "
                                    f"agent_model_config = :mconfig WHERE id = :id"
                                ),
                                {"mid": model_id_val, "mconfig": json.dumps(config_to_save), "id": row_id}
                            )
                    except Exception as e:
                        logging.error(f"Failed to migrate {table} ID {row_id}: {e}")
            except Exception as e:
                logging.error(f"Failed to query {table} for migration: {e}")

    # 3. 删除旧列
    for table in ['agent', 'agent_publish']:
        if column_exists(table, 'model'):
            with op.batch_alter_table(table, schema=None) as batch_op:
                batch_op.drop_column('model')
    # ### end Alembic commands ###


def downgrade() -> None:
    """Downgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    for table in ['agent_publish', 'agent']:
        with op.batch_alter_table(table, schema=None) as batch_op:
            if not column_exists(table, 'model'):
                batch_op.add_column(sa.Column('model', sqlite.JSON(), nullable=True))
            if column_exists(table, 'agent_model_config'):
                batch_op.drop_column('agent_model_config')
            if column_exists(table, 'model_id'):
                batch_op.drop_column('model_id')
    # ### end Alembic commands ###