"""删库重建脚本"""
import os
import sys
import pymysql
from pathlib import Path

SQL_DIR = Path(__file__).parent
SERVER_DIR = SQL_DIR.parent

# 加载 .env 文件
env_file = SERVER_DIR / '.env'
if env_file.exists():
    with open(env_file, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith('#') and '=' in line:
                key, value = line.split('=', 1)
                os.environ.setdefault(key.strip(), value.strip())

DB_CONFIG = {
    'host': os.getenv('DB_HOST', 'localhost'),
    'port': int(os.getenv('DB_PORT', 3306)),
    'user': os.getenv('DB_USER', 'root'),
    'password': os.getenv('DB_PASSWORD', '123456'),
    'charset': 'utf8mb4'
}

def read_sql_file(filename):
    with open(SQL_DIR / filename, 'r', encoding='utf-8') as f:
        return f.read()

def execute_sql(cursor, sql, desc):
    print(f'{desc}...')
    for stmt in [s.strip() for s in sql.split(';') if s.strip()]:
        try:
            cursor.execute(stmt)
        except pymysql.Error as e:
            if e.args[0] not in [1007, 1050]:
                print(f'  警告: {e.args[1]}')
    print(f'  {desc}完成！')

def rebuild():
    print('=' * 50)
    print('星域故事汇 - 删库重建')
    print('=' * 50)
    
    conn = pymysql.connect(**DB_CONFIG)
    cur = conn.cursor()
    
    # 删库
    print('删除旧数据库...')
    cur.execute('DROP DATABASE IF EXISTS stardom_story')
    conn.commit()
    print('  删除完成！')
    
    # 重建
    schema_sql = read_sql_file('schema.sql')
    execute_sql(cur, schema_sql, '创建数据库表结构')
    conn.commit()
    
    seed1 = read_sql_file('seed_stories_part1.sql')
    execute_sql(cur, seed1, '导入种子数据（第1部分）')
    conn.commit()
    
    seed2 = read_sql_file('seed_stories_part2.sql')
    execute_sql(cur, seed2, '导入种子数据（第2部分）')
    conn.commit()
    
    print('\n数据库重建完成！')
    cur.close()
    conn.close()

if __name__ == '__main__':
    rebuild()