#!/usr/bin/env python3
"""
重新初始化数据库表
用于在数据库表丢失时重新创建所有表结构
"""

import sys
import os
import asyncio

# 添加services/mcp-server到路径，以便导入模块
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'services', 'mcp-server'))

from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy import text
from models import Base
from config import settings
from database import prepare_database_url
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


async def reinit_database():
    """重新初始化数据库表"""
    try:
        # 准备数据库URL
        database_url = prepare_database_url(settings.database_url)
        
        # 确保使用asyncpg驱动
        if database_url.startswith("postgresql://") and "+asyncpg" not in database_url:
            database_url = database_url.replace("postgresql://", "postgresql+asyncpg://", 1)
        
        logger.info(f"正在连接到数据库: {database_url.split('@')[1] if '@' in database_url else '本地数据库'}")
        
        # 创建数据库引擎
        engine = create_async_engine(database_url, echo=False)
        
        logger.info("开始创建数据库表...")
        
        # 创建所有表
        async with engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)
        
        logger.info("✓ 数据库表创建成功！")
        
        # 列出所有创建的表
        async with engine.begin() as conn:
            # 获取所有表名
            if "postgresql" in database_url:
                result = await conn.execute(
                    text("""
                    SELECT table_name 
                    FROM information_schema.tables 
                    WHERE table_schema = 'public' 
                    ORDER BY table_name
                    """)
                )
            else:
                # SQLite
                result = await conn.execute(
                    text("""
                    SELECT name 
                    FROM sqlite_master 
                    WHERE type='table' AND name NOT LIKE 'sqlite_%'
                    ORDER BY name
                    """)
                )
            
            tables = result.fetchall()
            table_names = [row[0] for row in tables]
            
            logger.info(f"\n已创建的表列表（共 {len(table_names)} 个）:")
            for i, table_name in enumerate(table_names, 1):
                logger.info(f"  {i}. {table_name}")
        
        await engine.dispose()
        
        logger.info("\n数据库初始化完成！")
        logger.info("注意：此脚本只创建表结构，不会创建初始数据。")
        logger.info("如果需要创建初始数据，请运行服务或使用其他初始化脚本。")
        
        return True
        
    except Exception as e:
        logger.error(f"✗ 数据库初始化失败: {e}")
        import traceback
        traceback.print_exc()
        return False


async def main():
    """主函数"""
    print("="*80)
    print("重新初始化数据库表")
    print("="*80)
    print()
    
    # 确认操作
    print("警告：此操作将创建所有数据库表。")
    print("如果表已存在，SQLAlchemy不会删除或修改现有表。")
    print()
    
    success = await reinit_database()
    
    if success:
        print("\n" + "="*80)
        print("✓ 数据库表初始化成功！")
        print("="*80)
        sys.exit(0)
    else:
        print("\n" + "="*80)
        print("✗ 数据库表初始化失败！")
        print("="*80)
        sys.exit(1)


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