#!/usr/bin/env python3
"""
Test database connections for AI Analysis Service
"""

import os
import psycopg2
import redis
import pymongo
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

def test_postgres_connection():
    """Test PostgreSQL connection"""
    try:
        conn = psycopg2.connect(
            host=os.getenv('POSTGRES_HOST', 'localhost'),
            port=os.getenv('POSTGRES_PORT', 5432),
            database=os.getenv('POSTGRES_DB', 'dev_pipeline'),
            user=os.getenv('POSTGRES_USER', 'pipeline_admin'),
            password=os.getenv('POSTGRES_PASSWORD', 'secure_pipeline_2024')
        )
        
        cursor = conn.cursor()
        cursor.execute("SELECT COUNT(*) FROM all_repositories;")
        count = cursor.fetchone()[0]
        
        cursor.close()
        conn.close()
        
        print(f"✅ PostgreSQL: Connected successfully, {count} repositories found")
        return True
        
    except Exception as e:
        print(f"❌ PostgreSQL: Connection failed - {e}")
        return False

def test_redis_connection():
    """Test Redis connection"""
    try:
        r = redis.Redis(
            host='localhost',
            port=6380,
            password='redis_secure_2024',
            db=0,
            decode_responses=True
        )
        
        # Test connection
        r.ping()
        
        # Get database size
        dbsize = r.dbsize()
        
        print(f"✅ Redis: Connected successfully, {dbsize} keys found")
        return True
        
    except Exception as e:
        print(f"❌ Redis: Connection failed - {e}")
        return False

def test_mongodb_connection():
    """Test MongoDB connection"""
    try:
        client = pymongo.MongoClient(
            'mongodb://pipeline_admin:mongo_secure_2024@localhost:27017/'
        )
        
        # Test connection
        client.admin.command('ping')
        
        # Get database info
        db = client[os.getenv('MONGODB_DB', 'repo_analyzer')]
        collections = db.list_collection_names()
        
        print(f"✅ MongoDB: Connected successfully, {len(collections)} collections found")
        return True
        
    except Exception as e:
        print(f"❌ MongoDB: Connection failed - {e}")
        return False

def main():
    """Test all database connections"""
    print("🔍 Testing Database Connections...")
    print("=" * 50)
    
    postgres_ok = test_postgres_connection()
    redis_ok = test_redis_connection()
    mongodb_ok = test_mongodb_connection()
    
    print("=" * 50)
    print(f"📊 Connection Summary:")
    print(f"   PostgreSQL: {'✅' if postgres_ok else '❌'}")
    print(f"   Redis: {'✅' if redis_ok else '❌'}")
    print(f"   MongoDB: {'✅' if mongodb_ok else '❌'}")
    
    if all([postgres_ok, redis_ok, mongodb_ok]):
        print("🎉 All database connections successful!")
    else:
        print("⚠️  Some database connections failed")

if __name__ == "__main__":
    main()