#!/usr/bin/env python3
"""
数据库管理脚本
提供数据库初始化、检查、迁移等功能
"""

import asyncio
import sys
import os
import argparse

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(__file__))

from app.core.database import init_db, close_db, check_db_connection
from app.core.config import settings


async def check_connection():
    """检查数据库连接"""
    print(f"正在检查数据库连接...")
    print(f"数据库URL: {settings.DATABASE_URL}")
    
    is_connected = await check_db_connection()
    
    if is_connected:
        print("✅ 数据库连接正常!")
        return True
    else:
        print("❌ 数据库连接失败!")
        print("请检查:")
        print("1. PostgreSQL服务是否运行")
        print("2. 数据库配置是否正确")
        print("3. 网络连接是否正常")
        return False


async def initialize_database():
    """初始化数据库"""
    print(f"正在连接数据库: {settings.DATABASE_URL}")
    print("正在创建数据库表...")
    
    try:
        # 初始化数据库表
        await init_db()
        
        print("✅ 数据库表创建成功!")
        print("💡 提示: 如果需要使用Alembic管理迁移，请运行:")
        print("   alembic stamp head")
        print("   alembic revision --autogenerate -m '描述'")
        print("   alembic upgrade head")
        return True
        
    except Exception as e:
        print(f"❌ 数据库初始化失败: {e}")
        return False


async def show_status():
    """显示数据库状态"""
    print("=== 数据库状态 ===")
    print(f"数据库URL: {settings.DATABASE_URL}")
    print(f"数据库Echo: {settings.DATABASE_ECHO}")
    
    # 检查连接
    is_connected = await check_db_connection()
    print(f"连接状态: {'✅ 正常' if is_connected else '❌ 失败'}")
    
    if is_connected:
        try:
            from app.core.database import AsyncSessionLocal
            async with AsyncSessionLocal() as session:
                # 检查表是否存在
                result = await session.execute("""
                    SELECT table_name 
                    FROM information_schema.tables 
                    WHERE table_schema = 'public' 
                    AND table_name IN ('reports', 'analysis_modules', 'progress_tracking', 'system_config')
                    ORDER BY table_name
                """)
                tables = [row[0] for row in result.fetchall()]
                
                print(f"已创建的表: {', '.join(tables) if tables else '无'}")
                
                if 'reports' in tables:
                    result = await session.execute("SELECT COUNT(*) FROM reports")
                    count = result.scalar()
                    print(f"报告数量: {count}")
                    
        except Exception as e:
            print(f"获取详细状态失败: {e}")


async def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='数据库管理工具')
    parser.add_argument('command', choices=['check', 'init', 'status'], 
                       help='要执行的命令')
    
    args = parser.parse_args()
    
    try:
        if args.command == 'check':
            success = await check_connection()
        elif args.command == 'init':
            success = await initialize_database()
        elif args.command == 'status':
            await show_status()
            success = True
        else:
            print(f"未知命令: {args.command}")
            success = False
            
        if not success:
            sys.exit(1)
            
    except Exception as e:
        print(f"❌ 执行失败: {e}")
        sys.exit(1)
    finally:
        await close_db()


if __name__ == "__main__":
    asyncio.run(main())