#!/usr/bin/env python3
"""
PostgreSQL 数据库自动设置脚本

功能:
1. 自动连接到PostgreSQL服务器
2. 创建数据库和用户
3. 设置权限
4. 初始化表结构

使用方法:
    python backend/scripts/setup_postgres.py

前置条件:
    - PostgreSQL服务已安装并运行
    - 知道PostgreSQL的超级用户密码（通常是postgres用户）
"""
import sys
import asyncio
from pathlib import Path
from getpass import getpass
import logging

# 添加项目根目录到Python路径
sys.path.insert(0, str(Path(__file__).parent.parent))

try:
    import psycopg2
    from psycopg2 import sql
    from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
except ImportError:
    print("❌ 缺少psycopg2依赖，正在安装...")
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "psycopg2-binary"])
    import psycopg2
    from psycopg2 import sql
    from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT

# 注意: 表结构应由 Alembic 管理
from pathlib import Path

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(message)s'
)
logger = logging.getLogger(__name__)


class PostgreSQLSetup:
    """PostgreSQL数据库自动设置"""
    
    def __init__(
        self,
        host: str = "localhost",
        port: int = 5432,
        admin_user: str = "postgres",
        admin_password: str = None,
        db_name: str = "mumuai_novel",
        db_user: str = "mumuai",
        db_password: str = "123456"
    ):
        """
        初始化设置参数
        
        Args:
            host: PostgreSQL主机地址
            port: PostgreSQL端口
            admin_user: 管理员用户名
            admin_password: 管理员密码
            db_name: 要创建的数据库名
            db_user: 要创建的用户名
            db_password: 用户密码
        """
        self.host = host
        self.port = port
        self.admin_user = admin_user
        self.admin_password = admin_password
        self.db_name = db_name
        self.db_user = db_user
        self.db_password = db_password
        self.conn = None
    
    def connect_as_admin(self) -> bool:
        """连接到PostgreSQL（使用管理员权限）"""
        try:
            logger.info(f"🔌 连接到 PostgreSQL ({self.host}:{self.port})...")
            
            self.conn = psycopg2.connect(
                host=self.host,
                port=self.port,
                user=self.admin_user,
                password=self.admin_password,
                database="postgres"  # 连接到默认数据库
            )
            self.conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
            
            logger.info(f"✅ 已连接到 PostgreSQL")
            return True
            
        except psycopg2.OperationalError as e:
            logger.error(f"❌ 连接失败: {e}")
            logger.error("\n可能的原因:")
            logger.error("1. PostgreSQL服务未启动")
            logger.error("2. 管理员密码错误")
            logger.error("3. 主机地址或端口错误")
            logger.error("4. pg_hba.conf配置不允许连接")
            return False
    
    def database_exists(self) -> bool:
        """检查数据库是否存在"""
        cursor = self.conn.cursor()
        cursor.execute(
            "SELECT 1 FROM pg_database WHERE datname = %s",
            (self.db_name,)
        )
        exists = cursor.fetchone() is not None
        cursor.close()
        return exists
    
    def user_exists(self) -> bool:
        """检查用户是否存在"""
        cursor = self.conn.cursor()
        cursor.execute(
            "SELECT 1 FROM pg_user WHERE usename = %s",
            (self.db_user,)
        )
        exists = cursor.fetchone() is not None
        cursor.close()
        return exists
    
    def create_user(self) -> bool:
        """创建数据库用户"""
        try:
            if self.user_exists():
                logger.info(f"ℹ️  用户 '{self.db_user}' 已存在")
                
                # 询问是否重置密码
                response = input(f"是否重置用户 '{self.db_user}' 的密码? (yes/no): ")
                if response.lower() in ['yes', 'y']:
                    cursor = self.conn.cursor()
                    cursor.execute(
                        sql.SQL("ALTER USER {} WITH PASSWORD %s").format(
                            sql.Identifier(self.db_user)
                        ),
                        (self.db_password,)
                    )
                    cursor.close()
                    logger.info(f"✅ 用户密码已更新")
                
                return True
            
            logger.info(f"👤 创建用户 '{self.db_user}'...")
            cursor = self.conn.cursor()
            cursor.execute(
                sql.SQL("CREATE USER {} WITH PASSWORD %s").format(
                    sql.Identifier(self.db_user)
                ),
                (self.db_password,)
            )
            cursor.close()
            logger.info(f"✅ 用户创建成功")
            return True
            
        except Exception as e:
            logger.error(f"❌ 创建用户失败: {e}")
            return False
    
    def create_database(self) -> bool:
        """创建数据库"""
        try:
            if self.database_exists():
                logger.info(f"ℹ️  数据库 '{self.db_name}' 已存在")
                
                # 询问是否删除重建
                response = input(f"是否删除并重建数据库 '{self.db_name}'? (yes/no): ")
                if response.lower() in ['yes', 'y']:
                    logger.warning(f"⚠️  删除数据库 '{self.db_name}'...")
                    cursor = self.conn.cursor()
                    # 断开所有连接
                    cursor.execute(
                        sql.SQL("""
                            SELECT pg_terminate_backend(pg_stat_activity.pid)
                            FROM pg_stat_activity
                            WHERE pg_stat_activity.datname = %s
                            AND pid <> pg_backend_pid()
                        """),
                        (self.db_name,)
                    )
                    cursor.execute(
                        sql.SQL("DROP DATABASE {}").format(
                            sql.Identifier(self.db_name)
                        )
                    )
                    cursor.close()
                    logger.info(f"✅ 数据库已删除")
                else:
                    return True
            
            logger.info(f"🗄️  创建数据库 '{self.db_name}'...")
            cursor = self.conn.cursor()
            cursor.execute(
                sql.SQL("CREATE DATABASE {} OWNER {}").format(
                    sql.Identifier(self.db_name),
                    sql.Identifier(self.db_user)
                )
            )
            cursor.close()
            logger.info(f"✅ 数据库创建成功")
            return True
            
        except Exception as e:
            logger.error(f"❌ 创建数据库失败: {e}")
            return False
    
    def grant_privileges(self) -> bool:
        """授予用户权限"""
        try:
            logger.info(f"🔐 授予用户权限...")
            cursor = self.conn.cursor()
            
            # 授予数据库所有权限
            cursor.execute(
                sql.SQL("GRANT ALL PRIVILEGES ON DATABASE {} TO {}").format(
                    sql.Identifier(self.db_name),
                    sql.Identifier(self.db_user)
                )
            )
            
            cursor.close()
            logger.info(f"✅ 权限授予成功")
            return True
            
        except Exception as e:
            logger.error(f"❌ 授予权限失败: {e}")
            return False
    
    def update_env_file(self) -> bool:
        """更新.env文件"""
        try:
            env_file = Path(__file__).parent.parent / ".env"
            
            database_url = (
                f"postgresql+asyncpg://{self.db_user}:{self.db_password}"
                f"@{self.host}:{self.port}/{self.db_name}"
            )
            
            if env_file.exists():
                # 读取现有内容
                with open(env_file, 'r', encoding='utf-8') as f:
                    lines = f.readlines()
                
                # 更新DATABASE_URL
                updated = False
                for i, line in enumerate(lines):
                    if line.startswith('DATABASE_URL='):
                        lines[i] = f"DATABASE_URL={database_url}\n"
                        updated = True
                        break
                
                if not updated:
                    lines.append(f"\nDATABASE_URL={database_url}\n")
                
                # 写回文件
                with open(env_file, 'w', encoding='utf-8') as f:
                    f.writelines(lines)
            else:
                # 创建新文件
                with open(env_file, 'w', encoding='utf-8') as f:
                    f.write(f"DATABASE_URL={database_url}\n")
            
            logger.info(f"✅ .env 文件已更新")
            logger.info(f"   DATABASE_URL={database_url}")
            return True
            
        except Exception as e:
            logger.error(f"❌ 更新.env文件失败: {e}")
            return False
    
    async def initialize_tables(self) -> bool:
        """初始化数据库表结构（使用 Alembic）"""
        try:
            import subprocess
            logger.info(f"📋 使用 Alembic 初始化数据库表结构...")
            
            # 运行 Alembic 迁移
            result = subprocess.run(
                ["alembic", "upgrade", "head"],
                capture_output=True,
                text=True,
                cwd=Path(__file__).parent.parent
            )
            
            if result.returncode == 0:
                logger.info(f"✅ 表结构初始化成功")
                return True
            else:
                logger.error(f"❌ Alembic 迁移失败: {result.stderr}")
                return False
                
        except Exception as e:
            logger.error(f"❌ 初始化表结构失败: {e}")
            return False
    
    def close(self):
        """关闭数据库连接"""
        if self.conn:
            self.conn.close()
            logger.info(f"🔌 已断开连接")
    
    async def setup(self) -> bool:
        """执行完整设置流程"""
        try:
            # 1. 连接
            if not self.connect_as_admin():
                return False
            
            # 2. 创建用户
            if not self.create_user():
                return False
            
            # 3. 创建数据库
            if not self.create_database():
                return False
            
            # 4. 授予权限
            if not self.grant_privileges():
                return False
            
            # 5. 更新配置
            if not self.update_env_file():
                return False
            
            # 6. 关闭管理员连接
            self.close()
            
            # 7. 初始化表结构
            if not await self.initialize_tables():
                return False
            
            return True
            
        except Exception as e:
            logger.error(f"❌ 设置过程出错: {e}")
            return False
        finally:
            if self.conn:
                self.close()


async def main():
    """主函数"""
    print("""
╔═══════════════════════════════════════════════════════════════╗
║          PostgreSQL 数据库自动设置工具                        ║
║                                                               ║
║  此工具将自动完成:                                            ║
║  1. 连接到PostgreSQL服务器                                    ║
║  2. 创建数据库和用户                                          ║
║  3. 设置权限                                                  ║
║  4. 初始化表结构                                              ║
║  5. 更新.env配置文件                                          ║
╚═══════════════════════════════════════════════════════════════╝
    """)
    
    # 获取配置
    print("请输入PostgreSQL配置信息:\n")
    
    host = input("主机地址 [localhost]: ").strip() or "localhost"
    port = input("端口 [5432]: ").strip() or "5432"
    port = int(port)
    
    admin_user = input("管理员用户名 [postgres]: ").strip() or "postgres"
    admin_password = getpass(f"管理员密码: ")
    
    print("\n请输入要创建的数据库信息:\n")
    db_name = input("数据库名 [mumuai_novel]: ").strip() or "mumuai_novel"
    db_user = input("数据库用户名 [mumuai]: ").strip() or "mumuai"
    db_password = getpass("数据库用户密码 [mumuai123]: ") or "mumuai123"
    
    print(f"\n{'='*60}")
    print(f"配置摘要:")
    print(f"  服务器: {host}:{port}")
    print(f"  数据库: {db_name}")
    print(f"  用户: {db_user}")
    print(f"{'='*60}\n")
    
    response = input("确认开始设置? (yes/no): ")
    if response.lower() not in ['yes', 'y']:
        print("已取消设置")
        return
    
    # 执行设置
    setup = PostgreSQLSetup(
        host=host,
        port=port,
        admin_user=admin_user,
        admin_password=admin_password,
        db_name=db_name,
        db_user=db_user,
        db_password=db_password
    )
    
    print(f"\n{'='*60}")
    success = await setup.setup()
    print(f"{'='*60}\n")
    
    if success:
        print("🎉 PostgreSQL设置完成!\n")
        print("下一步:")
        print("1. 启动应用: python -m app.main")
        print("2. 访问: http://localhost:8000")
        print("3. 查看API文档: http://localhost:8000/docs")
    else:
        print("❌ 设置过程中出现错误，请检查日志")
        print("\n故障排查:")
        print("1. 确认PostgreSQL服务正在运行")
        print("2. 检查管理员用户名和密码")
        print("3. 查看PostgreSQL日志")


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