# -*- coding: utf-8 -*-
"""为 prompt_history 增加对比相关列（可重复执行，已存在则跳过）。"""
from __future__ import annotations

import os
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from dotenv import load_dotenv
from sqlalchemy import create_engine, text

load_dotenv(ROOT / ".env")

DATABASE_URL = os.environ.get("DATABASE_URL")
if not DATABASE_URL:
    print("ERROR: 未设置 DATABASE_URL", file=sys.stderr)
    sys.exit(1)

STATEMENTS = [
    (
        "comparison_group_id",
        "ALTER TABLE prompt_history ADD COLUMN comparison_group_id VARCHAR(64) NULL COMMENT '对比组ID'",
    ),
    (
        "is_comparison_enabled",
        "ALTER TABLE prompt_history ADD COLUMN is_comparison_enabled TINYINT(1) NOT NULL DEFAULT 0 COMMENT '是否参与对比'",
    ),
    (
        "idx_comparison_group",
        "ALTER TABLE prompt_history ADD INDEX idx_comparison_group (comparison_group_id)",
    ),
]


def main() -> int:
    engine = create_engine(DATABASE_URL)
    ok = 0
    with engine.connect() as conn:
        for name, sql in STATEMENTS:
            try:
                conn.execute(text(sql))
                conn.commit()
                print(f"OK: {name}")
                ok += 1
            except Exception as e:
                conn.rollback()
                err = str(e).lower()
                if "duplicate" in err or "1060" in str(e) or "1061" in str(e) or "1826" in str(e):
                    print(f"SKIP (已存在): {name}")
                else:
                    print(f"FAIL: {name} -> {e}", file=sys.stderr)
                    return 1
    print(f"完成，成功执行 {ok} 条新变更。")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())