# -*- coding: utf-8 -*-
"""迁移脚本：添加 token 消耗字段到对话记录表。"""
import logging
from sqlalchemy import text
from insurance.db.compat import db


def migrate():
    """添加 model_id、message_tokens、answer_tokens 字段到 insurance_chat_records 表。"""
    try:
        from sqlalchemy import inspect

        inspector = inspect(db.engine)
        columns = [col['name'] for col in inspector.get_columns('insurance_chat_records')]

        dialect = db.engine.dialect.name

        # 添加 model_id 字段
        if 'model_id' not in columns:
            if dialect == 'sqlite':
                db.session.execute(text(
                    "ALTER TABLE insurance_chat_records "
                    "ADD COLUMN model_id VARCHAR(100)"
                ))
            else:
                db.session.execute(text(
                    "ALTER TABLE insurance_chat_records "
                    "ADD COLUMN model_id VARCHAR(100) NULL "
                    "COMMENT '模型 ID'"
                ))
            logging.info("成功添加 model_id 字段")

        # 添加 message_tokens 字段（用户输入 token 数）
        if 'message_tokens' not in columns:
            if dialect == 'sqlite':
                db.session.execute(text(
                    "ALTER TABLE insurance_chat_records "
                    "ADD COLUMN message_tokens INTEGER DEFAULT 0"
                ))
            else:
                db.session.execute(text(
                    "ALTER TABLE insurance_chat_records "
                    "ADD COLUMN message_tokens INTEGER DEFAULT 0 "
                    "COMMENT '用户输入 token 数'"
                ))
            logging.info("成功添加 message_tokens 字段")

        # 添加 answer_tokens 字段（助手回复 token 数）
        if 'answer_tokens' not in columns:
            if dialect == 'sqlite':
                db.session.execute(text(
                    "ALTER TABLE insurance_chat_records "
                    "ADD COLUMN answer_tokens INTEGER DEFAULT 0"
                ))
            else:
                db.session.execute(text(
                    "ALTER TABLE insurance_chat_records "
                    "ADD COLUMN answer_tokens INTEGER DEFAULT 0 "
                    "COMMENT '助手回复 token 数'"
                ))
            logging.info("成功添加 answer_tokens 字段")

        db.session.commit()
        logging.info("迁移完成：token 字段添加成功")

    except Exception as e:
        logging.error("添加 token 字段失败: %s", e)
        db.session.rollback()
        raise


def downgrade():
    """删除 token 相关字段。"""
    try:
        from sqlalchemy import inspect

        inspector = inspect(db.engine)
        columns = [col['name'] for col in inspector.get_columns('insurance_chat_records')]

        for field in ['model_id', 'message_tokens', 'answer_tokens']:
            if field in columns:
                db.session.execute(text(f"ALTER TABLE insurance_chat_records DROP COLUMN {field}"))
                logging.info(f"成功删除 {field} 字段")

        db.session.commit()
        logging.info("回滚完成：token 字段删除成功")

    except Exception as e:
        logging.error("删除 token 字段失败: %s", e)
        db.session.rollback()
        raise


if __name__ == "__main__":
    import sys
    sys.path.insert(0, ".")
    from app import app
    with app.app_context():
        migrate()
        print("迁移完成")