#!/usr/bin/env python
"""数据库连接测试脚本"""
import sys
import os

# 添加项目根目录
project_root = os.path.abspath(os.path.dirname(__file__))
sys.path.insert(0, project_root)

try:
    import pg8000
    print("✅ pg8000 模块已安装")
except ImportError:
    print("❌ pg8000 模块未安装，请运行: pip install pg8000")
    sys.exit(1)

def test_connection(database_config):
    """测试数据库连接"""
    try:
        conn = pg8000.connect(
            host=database_config['host'],
            port=database_config['port'],
            user=database_config['user'],
            password=database_config['password'],
            database=database_config['database']
        )
        cursor = conn.cursor()
        cursor.execute("SELECT version();")
        version = cursor.fetchone()[0]
        print(f"✅ 连接成功! 数据库版本: {version[:50]}...")
        conn.close()
        return True
    except Exception as e:
        print(f"❌ 连接失败: {e}")
        return False

def main():
    print("=" * 60)
    print("数据库连接测试")
    print("=" * 60)
    print()

    # 从.env文件加载配置
    env_config = {}
    env_file = '.env'
    if os.path.exists(env_file):
        with open(env_file, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith('#') and '=' in line:
                    key, value = line.split('=', 1)
                    env_config[key] = value

    # 测试两个数据库
    databases = [
        {
            'name': 'fs_law_risk',
            'host': env_config.get('PG_HOST'),
            'port': int(env_config.get('PG_PORT', '5432')),
            'user': env_config.get('PG_USER'),
            'password': env_config.get('PG_PASSWORD'),
            'database': env_config.get('PG_DATABASE')
        },
        {
            'name': 'licensing_risks',
            'host': env_config.get('LIC_PG_HOST'),
            'port': int(env_config.get('LIC_PG_PORT', '5432')),
            'user': env_config.get('LIC_PG_USER'),
            'password': env_config.get('LIC_PG_PASSWORD'),
            'database': env_config.get('LIC_PG_DATABASE')
        }
    ]

    for db in databases:
        print(f"\n测试 {db['name']} 数据库...")
        print(f"  主机: {db['host']}:{db['port']}")
        print(f"  用户: {db['user']}")
        print(f"  数据库: {db['database']}")
        print(f"  密码: {'*' * len(db['password']) if db['password'] else 'N/A'}")

        if not all([db['host'], db['user'], db['password'], db['database']]):
            print("  ❌ 配置不完整")
            continue

        test_connection(db)

    print("\n" + "=" * 60)
    print("\n如果连接失败，请检查:")
    print("1. 数据库服务器是否运行")
    print("2. 网络是否可访问")
    print("3. 用户名/密码是否正确")
    print("4. 是否有防火墙阻止")
    print("=" * 60)

if __name__ == "__main__":
    main()