#!/usr/bin/env python3
"""
完整的数据库初始化脚本
包括创建表和插入测试数据
"""
import asyncio
import sys
from pathlib import Path
from sqlalchemy import text
from app.database import engine, Base, AsyncSessionLocal
from loguru import logger

# 导入初始化数据函数
from init_test_data import init_test_data

logger.remove()
logger.add(sys.stderr, level="INFO", format="<green>{time:HH:mm:ss}</green> | <level>{level: <8}</level> | {message}")


async def test_database_connection():
    """测试数据库连接"""
    logger.info("正在测试数据库连接...")

    try:
        async with engine.connect() as conn:
            result = await conn.execute(text("SELECT version()"))
            version = result.fetchone()
            logger.info(f"✅ 数据库连接成功")
            logger.info(f"数据库版本: {version[0] if version else 'N/A'}")
            return True
    except Exception as e:
        logger.error(f"❌ 数据库连接失败: {e}")
        logger.error("\n请检查:")
        logger.error("1. PostgreSQL 服务是否运行")
        logger.error("2. 数据库 'deeprisk_db' 是否存在")
        logger.error("3. 用户名密码是否正确")
        logger.error(f"\n当前数据库URL配置:")
        logger.error(f"postgresql+asyncpg://postgres:***@localhost:15432/deeprisk_db")
        return False


async def init_tables():
    """初始化数据库表"""
    logger.info("=" * 60)
    logger.info("正在初始化数据库表结构...")
    logger.info("=" * 60)

    try:
        async with engine.begin() as conn:
            # 创建所有表
            await conn.run_sync(Base.metadata.create_all)
            logger.info("✅ 数据库表创建成功")

            # 列出已创建的表
            result = await conn.execute(
                text("""
                    SELECT table_name
                    FROM information_schema.tables
                    WHERE table_schema = 'public'
                    ORDER BY table_name
                """)
            )
            tables = result.fetchall()
            if tables:
                logger.info("\n已创建的表:")
                for table in tables:
                    logger.info(f"  ✓ {table[0]}")

    except Exception as e:
        logger.error(f"❌ 创建数据库表失败: {e}")
        import traceback
        traceback.print_exc()
        raise


async def show_data_summary():
    """显示数据统计"""
    logger.info("\n" + "=" * 60)
    logger.info("数据库数据统计")
    logger.info("=" * 60)

    async with AsyncSessionLocal() as session:
        try:
            # 检查各表的数据量
            tables = [
                "mcn_agency",
                "streamer_info",
                "platform_recharge",
                "contract",
                "order",
                "settlement",
                "expense",
                "tax_declaration",
                "bank_transaction",
                "invoice",
                "detection_rule",
            ]

            for table in tables:
                try:
                    result = await session.execute(text(f"SELECT COUNT(*) FROM {table}"))
                    count = result.scalar()
                    logger.info(f"  {table:30} : {count:6} 条记录")
                except Exception as e:
                    logger.warning(f"  {table:30} : 查询失败 ({str(e)[:30]})")

        except Exception as e:
            logger.error(f"统计查询失败: {e}")


async def main():
    """主函数"""
    logger.info("\n" + "=" * 60)
    logger.info("🎯 DeepRisk 数据库完整初始化")
    logger.info("=" * 60 + "\n")

    # 1. 测试数据库连接
    if not await test_database_connection():
        sys.exit(1)

    # 2. 初始化表结构
    await init_tables()

    # 3. 初始化测试数据
    try:
        await init_test_data()
    except Exception as e:
        logger.error(f"初始化测试数据失败: {e}")
        logger.warning("表创建成功，但数据插入失败")
        logger.warning("可以稍后手动运行 init_test_data.py 插入数据")

    # 4. 显示数据统计
    await show_data_summary()

    logger.info("\n" + "=" * 60)
    logger.info("🎉 数据库初始化完成！")
    logger.info("=" * 60)
    logger.info("\n可以使用以下方式启动应用:")
    logger.info("  uvicorn app.main:app --reload --host 0.0.0.0 --port 8000\n")

    await engine.dispose()


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