#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据库初始化脚本
用于创建所有必需的数据库表，确保数据库结构与模型定义一致

使用方法：
    python3 scripts/init_db.py                      # 正常初始化（保留现有数据）
    python3 scripts/init_db.py --reset              # 完全重置数据库（删除所有数据）
    python3 scripts/init_db.py --reset -u admin -e admin@example.com -p password123
                                                    # 重置数据库并创建管理员用户
"""

import os
import sys
import argparse
import re

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

from flask import Flask
from app.models.database.models import db, User
from config import Config


def validate_email(email):
    """验证邮箱格式"""
    pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$'
    return re.match(pattern, email) is not None


def validate_username(username):
    """验证用户名格式"""
    if len(username) < 3 or len(username) > 20:
        return False, "用户名长度必须在3-20个字符之间"
    if not re.match(r'^[a-zA-Z0-9_]+$', username):
        return False, "用户名只能包含字母、数字和下划线"
    return True, ""


def validate_password(password):
    """验证密码强度"""
    if len(password) < 6:
        return False, "密码长度至少6个字符"
    return True, ""


def create_admin_user(username, email, password, skip_validation=False):
    """创建管理员用户
    
    Args:
        username: 用户名
        email: 邮箱
        password: 密码
        skip_validation: 是否跳过验证（用于命令行参数）
    """
    if not skip_validation:
        # 验证用户名
        valid, message = validate_username(username)
        if not valid:
            print(f"❌ {message}")
            return False
        
        # 验证邮箱
        if not validate_email(email):
            print("❌ 邮箱格式不正确")
            return False
        
        # 验证密码
        valid, message = validate_password(password)
        if not valid:
            print(f"❌ {message}")
            return False
    
    # 检查用户名是否已存在
    existing = User.query.filter_by(username=username).first()
    if existing:
        print(f"⚠️  用户名 '{username}' 已存在，将更新为管理员")
        if existing.role != 'admin':
            existing.role = 'admin'
            existing.set_password(password)
            existing.email = email
            existing.is_active = True
            db.session.commit()
            print(f"✅ 已更新用户 '{username}' 为管理员")
        return True
    
    # 检查邮箱是否已存在
    existing = User.query.filter_by(email=email).first()
    if existing:
        print(f"⚠️  邮箱 '{email}' 已被使用，将更新为管理员")
        if existing.role != 'admin':
            existing.role = 'admin'
            existing.set_password(password)
            existing.username = username
            existing.is_active = True
            db.session.commit()
            print(f"✅ 已更新用户为管理员")
        return True
    
    # 创建新管理员用户
    try:
        admin = User(
            username=username,
            email=email,
            role='admin',
            is_active=True
        )
        admin.set_password(password)
        
        db.session.add(admin)
        db.session.commit()
        
        print(f"✅ 管理员账户创建成功！")
        print(f"   用户名: {username}")
        print(f"   邮箱: {email}")
        return True
        
    except Exception as e:
        db.session.rollback()
        print(f"❌ 创建管理员账户失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return False


def init_default_system_configs():
    """初始化默认的系统配置"""
    from app.models.database.models import SystemConfig

    print("\n开始初始化默认系统配置...")

    # 定义默认配置
    default_configs = [
        {
            'key': 'allow_registration',
            'value': True,
            'description': '是否允许用户注册（关闭后只能由管理员创建用户）',
            'category': 'security'
        }
    ]

    # 插入默认配置
    created_count = 0
    skipped_count = 0

    for config_data in default_configs:
        # 检查配置是否已存在
        existing = SystemConfig.query.filter_by(key=config_data['key']).first()
        if existing:
            print(f"  ⊙ 配置已存在，跳过: {config_data['key']}")
            skipped_count += 1
            continue

        # 创建新配置
        config = SystemConfig(
            key=config_data['key'],
            value='',  # 临时值，稍后通过set_value设置
            description=config_data['description'],
            category=config_data['category'],
            enabled=True
        )
        config.set_value(config_data['value'])
        db.session.add(config)
        print(f"  ✓ 创建配置: {config_data['key']} = {config_data['value']}")
        created_count += 1

    # 提交到数据库
    db.session.commit()

    print(f"\n系统配置初始化完成！")
    print(f"  - 新创建: {created_count} 个")
    print(f"  - 已存在: {skipped_count} 个")

    return True


def init_database(reset=False, admin_username=None, admin_email=None, admin_password=None, skip_confirm=False):
    """初始化数据库，创建所有表

    Args:
        reset: 是否完全重置数据库（删除所有表后重建）
        admin_username: 管理员用户名（可选）
        admin_email: 管理员邮箱（可选）
        admin_password: 管理员密码（可选）
        skip_confirm: 是否跳过确认（用于命令行参数）
    """
    # 创建Flask应用
    app = Flask(__name__)
    app.config.from_object(Config)

    # 数据库配置已经在 Config 中设置好了
    # 直接使用 Config 中的 SQLALCHEMY_DATABASE_URI
    print(f"\n数据库类型: {Config.DATABASE_TYPE}")
    print(f"数据库连接: {Config.SQLALCHEMY_DATABASE_URI}")

    # 初始化数据库
    db.init_app(app)

    with app.app_context():
        # 导入所有模型，确保它们被注册
        from app.models.database.models import (
            # 用户相关
            User, UserSession,
            # 题目相关
            ChallengeRecord,
            # 部署相关
            DeploymentRecord,
            # 优化记录
            OptimizeRecord,
            # 系统配置
            SystemConfig,
            # Augment配置
            AugmentSessionToken,
            # AnyRouter/AgentRouter Token
            AnyRouterToken, AgentRouterToken,
            # 方向配置相关
            CategoryConfig, KnowledgeCategory, KnowledgeItem,
            WriteupFile,
            CategoryAdmin, PendingWriteup,
            # 通知相关
            Notification
        )
        # 导入 AI 提供商配置（这些不在 models.py 兼容层，需要单独导入）
        from app.models.database import AIProviderConfig, SystemAIConfig

        if reset:
            if not skip_confirm:
                print("\n⚠️  警告：即将删除所有数据库表和数据！")
                confirm = input("请输入 'YES' 确认重置数据库: ")
                if confirm != 'YES':
                    print("❌ 操作已取消")
                    return False
            else:
                print("\n⚠️  警告：正在删除所有数据库表和数据...")

            print("\n正在删除所有表...")
            
            # 对于 PostgreSQL，需要处理外键依赖
            # 使用 CASCADE 删除所有表以处理外键依赖
            if Config.DATABASE_TYPE == 'postgresql':
                try:
                    # 获取所有表名
                    from sqlalchemy import inspect, text
                    inspector = inspect(db.engine)
                    tables = inspector.get_table_names()
                    
                    if tables:
                        print(f"  发现 {len(tables)} 个表，使用 CASCADE 删除...")
                        # 使用事务来确保所有删除操作原子性
                        with db.engine.begin() as conn:
                            # 直接删除所有表，使用 CASCADE 处理依赖
                            for table in tables:
                                try:
                                    conn.execute(text(f'DROP TABLE IF EXISTS "{table}" CASCADE'))
                                    print(f"  ✓ 已删除表: {table}")
                                except Exception as e:
                                    print(f"  ⚠ 删除表 {table} 时出错: {str(e)}")
                        print("✓ 所有表已删除（使用 CASCADE）")
                    else:
                        print("  ⊙ 数据库中没有表需要删除")
                except Exception as e:
                    print(f"  ⚠ CASCADE 删除失败，尝试标准删除: {str(e)}")
                    import traceback
                    traceback.print_exc()
                    # 回退到标准删除
                    try:
                        db.drop_all()
                        print("✓ 所有表已删除（标准方式）")
                    except Exception as e2:
                        print(f"  ✗ 标准删除也失败: {str(e2)}")
                        raise
            else:
                # SQLite 或其他数据库使用标准方式
                db.drop_all()
                print("✓ 所有表已删除")

        print("\n开始创建数据库表...")

        # 创建所有表
        db.create_all()

        print("✓ 数据库表创建完成！")
        if Config.DATABASE_TYPE == 'sqlite':
            print(f"数据库文件位置: {Config.DB_PATH}")
        else:
            print(f"PostgreSQL 数据库: {Config.POSTGRES_DB}")

        # 显示创建的表
        from sqlalchemy import inspect
        inspector = inspect(db.engine)
        tables = inspector.get_table_names()

        print(f"\n已创建 {len(tables)} 个表:")
        for table in sorted(tables):
            print(f"  - {table}")

        # 初始化默认系统配置
        init_default_system_configs()

        # 如果提供了管理员信息，创建管理员用户
        if admin_username and admin_email and admin_password:
            print("\n" + "=" * 60)
            print("创建管理员账户")
            print("=" * 60)
            create_admin_user(admin_username, admin_email, admin_password, skip_validation=True)

        return True


if __name__ == "__main__":
    # 解析命令行参数
    parser = argparse.ArgumentParser(
        description='数据库初始化脚本',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  python3 scripts/init_db.py                      # 正常初始化（保留现有数据）
  python3 scripts/init_db.py --reset              # 完全重置数据库（删除所有数据）
  python3 scripts/init_db.py --reset -u admin -e admin@example.com -p password123
                                                  # 重置数据库并创建管理员用户
        """
    )
    parser.add_argument(
        '--reset', '--force',
        action='store_true',
        dest='reset',
        help='完全重置数据库（删除所有表和数据后重建）'
    )
    parser.add_argument(
        '-u', '--username',
        dest='username',
        help='管理员用户名（与 --reset 一起使用时自动创建管理员）'
    )
    parser.add_argument(
        '-e', '--email',
        dest='email',
        help='管理员邮箱（与 --reset 一起使用时自动创建管理员）'
    )
    parser.add_argument(
        '-p', '--password',
        dest='password',
        help='管理员密码（与 --reset 一起使用时自动创建管理员）'
    )
    parser.add_argument(
        '--yes',
        action='store_true',
        dest='skip_confirm',
        help='跳过确认提示（与 --reset 一起使用）'
    )

    args = parser.parse_args()

    # 如果提供了管理员信息，检查是否完整
    admin_info_provided = any([args.username, args.email, args.password])
    if admin_info_provided:
        if not all([args.username, args.email, args.password]):
            print("❌ 错误：如果提供管理员信息，必须同时提供用户名(-u)、邮箱(-e)和密码(-p)")
            sys.exit(1)
        if not args.reset:
            print("❌ 错误：创建管理员用户只能在重置数据库时进行")
            print("   请使用: python3 scripts/init_db.py --reset -u <username> -e <email> -p <password>")
            sys.exit(1)

    try:
        if args.reset:
            print("\n" + "=" * 60)
            print("数据库完全重置模式")
            print("=" * 60)
            if admin_info_provided:
                print(f"将创建管理员用户: {args.username} ({args.email})")
        else:
            print("\n" + "=" * 60)
            print("数据库初始化")
            print("=" * 60)

        success = init_database(
            reset=args.reset,
            admin_username=args.username,
            admin_email=args.email,
            admin_password=args.password,
            skip_confirm=args.skip_confirm or (args.reset and admin_info_provided)
        )

        if success:
            print("\n" + "=" * 60)
            print("✓ 数据库初始化成功！")
            print("=" * 60)
            if args.reset and admin_info_provided:
                print(f"\n管理员账户已创建，可以使用以下信息登录：")
                print(f"  用户名: {args.username}")
                print(f"  邮箱: {args.email}")
            sys.exit(0)
        else:
            print("\n" + "=" * 60)
            print("✗ 数据库初始化失败！")
            print("=" * 60)
            sys.exit(1)
    except Exception as e:
        print(f"\n✗ 数据库初始化出错: {str(e)}")
        import traceback
        traceback.print_exc()
        sys.exit(1)