#!/usr/bin/env python3
"""
测试PostgreSQL数据库连接
"""

import os
import psycopg2
from psycopg2.extras import RealDictCursor
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

def test_connection():
    """测试数据库连接"""
    try:
        # 从环境变量读取配置
        conn_params = {
            'host': os.getenv('KNOWHUB_DB'),
            'port': int(os.getenv('KNOWHUB_PORT', 5432)),
            'user': os.getenv('KNOWHUB_USER'),
            'password': os.getenv('KNOWHUB_PASSWORD'),
            'database': os.getenv('KNOWHUB_DB_NAME'),
        }

        print("正在连接到PostgreSQL...")
        print(f"主机: {conn_params['host']}")
        print(f"端口: {conn_params['port']}")
        print(f"用户: {conn_params['user']}")
        print(f"数据库: {conn_params['database']}")
        print("-" * 50)

        # 建立连接
        conn = psycopg2.connect(**conn_params)
        cursor = conn.cursor(cursor_factory=RealDictCursor)

        # 测试查询
        cursor.execute("SELECT version();")
        version = cursor.fetchone()
        print(f"✅ 连接成功！")
        print(f"PostgreSQL版本: {version['version']}")
        print("-" * 50)

        # 检查pgvector扩展
        cursor.execute("""
            SELECT * FROM pg_extension WHERE extname = 'vector';
        """)
        pgvector = cursor.fetchone()
        if pgvector:
            print(f"✅ pgvector扩展已安装: {pgvector['extversion']}")
        else:
            print("⚠️  pgvector扩展未安装，需要先安装才能存储向量")
            print("   安装命令: CREATE EXTENSION vector;")
        print("-" * 50)

        # 列出所有表
        cursor.execute("""
            SELECT tablename FROM pg_tables
            WHERE schemaname = 'public'
            ORDER BY tablename;
        """)
        tables = cursor.fetchall()
        if tables:
            print(f"数据库中的表 ({len(tables)}个):")
            for table in tables:
                print(f"  - {table['tablename']}")
        else:
            print("数据库中暂无表")

        cursor.close()
        conn.close()
        print("-" * 50)
        print("✅ 测试完成，连接正常！")
        return True

    except Exception as e:
        print(f"❌ 连接失败: {e}")
        return False


if __name__ == "__main__":
    test_connection()