"""agent model

Revision ID: 9fa66e01df94
Revises: 06a1f79bce8b
Create Date: 2026-01-31 14:44:59.028934

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

from alembic import op
import sqlalchemy as sa
from sqlalchemy import inspect
from sqlalchemy.dialects import mysql

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

# revision identifiers, used by Alembic.
revision: str = '9fa66e01df94'
down_revision: Union[str, Sequence[str], None] = '06a1f79bce8b'
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. 为 agent 和 agent_publish 表添加新列
    for table in ['agent', 'agent_publish']:
        if not column_exists(table, 'model_id'):
            op.add_column(table, sa.Column('model_id', sa.Integer(), nullable=True))
        if not column_exists(table, 'agent_model_config'):
            op.add_column(table, sa.Column('agent_model_config', sa.JSON(), nullable=True))
    
    # 2. 数据迁移：从旧的 model JSON 字段提取 model_id 和 model_info
    conn = op.get_bind()
    for table in ['agent', 'agent_publish']:
        if column_exists(table, 'model'):
            # 查询现有数据
            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]
                        
                        # 更新 model_id 和 agent_model_config
                        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}")

    # 3. 删除旧的 model 列
    if column_exists('agent', 'model'):
        op.drop_column('agent', 'model')
    if column_exists('agent_publish', 'model'):
        op.drop_column('agent_publish', 'model')
    # ### end Alembic commands ###


def downgrade() -> None:
    """Downgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    # 1. 恢复 model 列
    for table in ['agent_publish', 'agent']:
        if not column_exists(table, 'model'):
            op.add_column(table, sa.Column('model', mysql.JSON(), nullable=True))
    
    # 2. 删除新添加的列
    for table in ['agent', 'agent_publish']:
        if column_exists(table, 'agent_model_config'):
            op.drop_column(table, 'agent_model_config')
        if column_exists(table, 'model_id'):
            op.drop_column(table, 'model_id')
    # ### end Alembic commands ###
