"""添加code字段的迁移脚本"""
import sys
import os
import random

# 添加项目根目录到系统路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app import create_app
from models import db, Site
from sqlalchemy import text

# 创建应用上下文
app = create_app()

def generate_unique_code():
    """生成唯一的8位数字编码"""
    while True:
        # 生成10000000-99999999之间的随机数
        code = str(random.randint(10000000, 99999999))
        # 检查是否已存在
        if not Site.query.filter_by(code=code).first():
            return code

with app.app_context():
    try:
        print("Step 1: Adding code column to sites table...")

        # 添加code列（先允许为空）
        with db.engine.connect() as connection:
            connection.execute(text('ALTER TABLE sites ADD COLUMN code VARCHAR(8) NULL COMMENT "8位数字编码"'))
            connection.commit()

        print("  - Column added successfully!")

        print("\nStep 2: Generating codes for existing sites...")

        # 为现有网站生成code
        sites = Site.query.all()
        for site in sites:
            site.code = generate_unique_code()
            print(f"  - Site #{site.id} '{site.name}': code = {site.code}")

        db.session.commit()
        print(f"  - Generated codes for {len(sites)} sites!")

        print("\nStep 3: Making code column NOT NULL and UNIQUE...")

        # 现在修改列为NOT NULL和UNIQUE
        with db.engine.connect() as connection:
            connection.execute(text('ALTER TABLE sites MODIFY COLUMN code VARCHAR(8) NOT NULL'))
            connection.execute(text('ALTER TABLE sites ADD UNIQUE INDEX idx_site_code (code)'))
            connection.commit()

        print("  - Code column constraints added!")
        print("\n✓ Migration completed successfully!")

    except Exception as e:
        print(f"\n✗ Migration failed: {str(e)}")
        import traceback
        traceback.print_exc()
        db.session.rollback()