#!/usr/bin/env python3
"""测试各个数据库连接和查询"""

import os
import json
import psycopg2
import pymysql
import requests
from requests.auth import HTTPBasicAuth
import warnings
warnings.filterwarnings('ignore')

def test_postgresql():
    """测试PostgreSQL连接"""
    print("\n" + "="*60)
    print("测试 PostgreSQL（Online）连接")
    print("="*60)
    
    try:
        conn = psycopg2.connect(
            host="bj-postgres-16pob4sg.sql.tencentcdb.com",
            port=28591,
            user="ai_member",
            password="LdfjdjL83h3h3^$&**YGG*",
            dbname="vala",
            connect_timeout=10
        )
        print("✅ PostgreSQL 连接成功！")
        
        # 测试查询
        with conn.cursor() as cur:
            # 先查询所有表
            cur.execute("SELECT tablename FROM pg_tables WHERE schemaname = 'public' LIMIT 5")
            tables = cur.fetchall()
            print(f"✅ 查询成功！找到前5个表：{[t[0] for t in tables]}")
            
            # 尝试查询其中一个表的1条数据
            if tables:
                table = tables[0][0]
                cur.execute(f"SELECT * FROM {table} LIMIT 1")
                row = cur.fetchone()
                print(f"✅ 从表 {table} 读取到1条数据：{row if row else '空表'}")
        
        conn.close()
        return True
        
    except Exception as e:
        print(f"❌ PostgreSQL 连接/查询失败：{str(e)[:200]}")
        return False

def test_mysql_test():
    """测试Test MySQL连接"""
    print("\n" + "="*60)
    print("测试 MySQL（Test环境）连接")
    print("="*60)
    
    try:
        conn = pymysql.connect(
            host="bj-cdb-8frbdwju.sql.tencentcdb.com",
            port=25413,
            user="read_only",
            password="fdsfiidier^$*hjfdijjd232",
            connect_timeout=10
        )
        print("✅ MySQL（Test）连接成功！")
        
        # 测试查询
        with conn.cursor() as cur:
            cur.execute("SHOW DATABASES LIMIT 5")
            dbs = cur.fetchall()
            print(f"✅ 查询成功！找到前5个数据库：{[db[0] for db in dbs]}")
            
            if dbs:
                db = dbs[0][0]
                cur.execute(f"USE {db}")
                cur.execute("SHOW TABLES LIMIT 1")
                table = cur.fetchone()
                if table:
                    cur.execute(f"SELECT * FROM {table[0]} LIMIT 1")
                    row = cur.fetchone()
                    print(f"✅ 从表 {table[0]} 读取到1条数据：{row if row else '空表'}")
        
        conn.close()
        return True
        
    except Exception as e:
        print(f"❌ MySQL（Test）连接/查询失败：{str(e)[:200]}")
        return False

def test_mysql_online():
    """测试Online MySQL连接"""
    print("\n" + "="*60)
    print("测试 MySQL（Online）连接")
    print("="*60)
    
    try:
        conn = pymysql.connect(
            host="bj-cdb-dh2fkqa0.sql.tencentcdb.com",
            port=27751,
            user="read_only",
            password="fsdo45ijfmfmuu77$%^&",
            connect_timeout=10
        )
        print("✅ MySQL（Online）连接成功！")
        
        # 测试查询
        with conn.cursor() as cur:
            cur.execute("SHOW DATABASES LIMIT 5")
            dbs = cur.fetchall()
            print(f"✅ 查询成功！找到前5个数据库：{[db[0] for db in dbs]}")
            
        conn.close()
        return True
        
    except Exception as e:
        print(f"❌ MySQL（Online）连接/查询失败：{str(e)[:200]}")
        return False

def test_es_online():
    """测试Online ES连接"""
    print("\n" + "="*60)
    print("测试 Elasticsearch（Online）连接")
    print("="*60)
    
    try:
        url = "https://es-7vd7jcu9.public.tencentelasticsearch.com:9200"
        auth = HTTPBasicAuth("elastic", "F%?QDcWes7N2WTuiYD11")
        
        response = requests.get(
            url,
            auth=auth,
            verify=False,
            timeout=10
        )
        
        if response.status_code == 200:
            info = response.json()
            print(f"✅ ES 连接成功！集群名称：{info.get('cluster_name')}")
            
            # 测试查询索引
            indices_resp = requests.get(
                f"{url}/_cat/indices?format=json",
                auth=auth,
                verify=False,
                timeout=10
            )
            if indices_resp.status_code == 200:
                indices = indices_resp.json()
                print(f"✅ 查询成功！索引数量：{len(indices)}")
                if indices:
                    print(f"   前3个索引：{[idx['index'] for idx in indices[:3]]}")
            
            return True
        else:
            print(f"❌ ES 连接失败：HTTP {response.status_code}")
            return False
            
    except Exception as e:
        print(f"❌ ES 连接/查询失败：{str(e)[:200]}")
        return False

if __name__ == "__main__":
    print("开始测试所有数据库连接...")
    
    results = {}
    results["PostgreSQL(Online)"] = test_postgresql()
    results["MySQL(Test)"] = test_mysql_test()
    results["MySQL(Online)"] = test_mysql_online()
    results["ES(Online)"] = test_es_online()
    
    print("\n" + "="*60)
    print("测试总结")
    print("="*60)
    for name, result in results.items():
        status = "✅ 正常" if result else "❌ 异常"
        print(f"{name}: {status}")