#!/usr/bin/env python3
"""只导入用户迁移所需的表到临时数据库 xaio_old"""
import MySQLdb
import re
from app_secrets import DATABASE_HOST, DATABASE_PORT, DATABASE_USER, DATABASE_PASSWORD

# 需要导入的表（用户相关）
NEEDED_TABLES = {
    'user_main', 'user_dashou', 'user_guanshi', 'user_kefu',
    'user_shangjia', 'user_boss', 'admin_profile', 'user_zuzhang',
    'user_shenheguan', 'user_role', 'role', 'permission',
    'role_permission', 'account_permission',
}

def get_conn(db=None):
    kwargs = dict(
        host=DATABASE_HOST, port=int(DATABASE_PORT),
        user=DATABASE_USER, password=DATABASE_PASSWORD,
        charset='utf8mb4', connect_timeout=60,
        read_timeout=300, write_timeout=300,
    )
    if db:
        kwargs['database'] = db
    return MySQLdb.connect(**kwargs)

# Drop and recreate
conn = get_conn()
cursor = conn.cursor()
cursor.execute('DROP DATABASE IF EXISTS xaio_old')
cursor.execute('CREATE DATABASE xaio_old CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci')
conn.commit()
conn.close()
print('临时数据库 xaio_old 已重建')

# Read the SQL file
with open('xaio_cheng_xu_backup.sql', 'r', encoding='utf-8') as f:
    sql_content = f.read()

# Parse and filter statements
# We need CREATE TABLE and INSERT statements for NEEDED_TABLES only
# Also need DROP TABLE IF EXISTS for those tables
statements = []
current = []
current_table = None
in_create = False

for line in sql_content.split('\n'):
    stripped = line.strip()
    if not stripped or stripped.startswith('--') or stripped.startswith('/*'):
        continue

    # Detect table name from CREATE TABLE or DROP TABLE or INSERT INTO
    create_match = re.match(r'CREATE\s+TABLE\s+`(\w+)`', stripped, re.IGNORECASE)
    drop_match = re.match(r'DROP\s+TABLE\s+IF\s+EXISTS\s+`(\w+)`', stripped, re.IGNORECASE)
    insert_match = re.match(r'INSERT\s+INTO\s+`(\w+)`', stripped, re.IGNORECASE)

    if create_match:
        current_table = create_match.group(1)
        in_create = True
    elif drop_match:
        current_table = drop_match.group(1)
    elif insert_match:
        current_table = insert_match.group(1)
        in_create = False

    # Only include statements for needed tables
    if current_table and current_table in NEEDED_TABLES:
        current.append(line)

        if stripped.endswith(';'):
            stmt = '\n'.join(current).strip()
            current = []
            current_table = None
            in_create = False
            upper = stmt.upper()
            if not upper.startswith('SET ') and not upper.startswith('USE '):
                statements.append(stmt)
    elif stripped.endswith(';'):
        # Not a needed table, skip
        current = []
        current_table = None
        in_create = False
    elif current_table and current_table in NEEDED_TABLES:
        current.append(line)
    else:
        # Not needed, skip but track state
        if stripped.endswith(';'):
            current = []
            current_table = None

print(f'筛选出 {len(statements)} 条用户相关SQL语句')

# Execute
conn = get_conn(db='xaio_old')
cursor = conn.cursor()
cursor.execute('SET FOREIGN_KEY_CHECKS = 0')
executed = 0
errors = 0

for stmt in statements:
    try:
        cursor.execute(stmt)
        executed += 1
        if executed % 20 == 0:
            conn.commit()
    except MySQLdb.OperationalError as e:
        if e.args[0] in (2006, 2013):
            try:
                conn.close()
            except:
                pass
            conn = get_conn(db='xaio_old')
            cursor = conn.cursor()
            try:
                cursor.execute(stmt)
                executed += 1
            except Exception as e2:
                errors += 1
                if errors <= 5:
                    print(f'  ERR(retry): {str(e2)[:100]}')
        else:
            errors += 1
            if errors <= 5:
                print(f'  ERR: {str(e)[:100]}')
    except Exception as e:
        errors += 1
        if errors <= 5:
            print(f'  ERR: {str(e)[:100]}')

conn.commit()
cursor.execute('SET FOREIGN_KEY_CHECKS = 1')
conn.commit()
print(f'\n导入完成: 成功 {executed}, 错误 {errors}')

# Check row counts
print('\n关键表行数:')
for table in sorted(NEEDED_TABLES):
    try:
        cursor.execute('SELECT COUNT(*) FROM `%s`' % table)
        count = cursor.fetchone()[0]
        print(f'  {table}: {count} rows')
    except Exception:
        print(f'  {table}: NOT FOUND')

conn.close()