"""
Database connection and utilities for MongoDB integration
"""

import os
from pymongo import MongoClient
from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError
import logging
from typing import Optional
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Configure logging
logger = logging.getLogger(__name__)

class DatabaseConnection:
    """MongoDB connection manager"""
    
    def __init__(self, mongodb_uri=None, database_name=None):
        self.client: Optional[MongoClient] = None
        self.db = None
        self.mongodb_url = mongodb_uri or os.getenv('MONGODB_URL') or os.getenv('MONGODB_URI')
        self.database_name = database_name or os.getenv('MONGODB_DATABASE', 'Atlas')
        
        if not self.mongodb_url:
            raise ValueError("MONGODB_URL environment variable is required")
    
    def connect(self):
        """Establish connection to MongoDB"""
        try:
            self.client = MongoClient(
                self.mongodb_url,
                serverSelectionTimeoutMS=5000,  # 5 second timeout
                connectTimeoutMS=10000,  # 10 second connection timeout
                socketTimeoutMS=20000,   # 20 second socket timeout
                maxPoolSize=50,          # Maximum connection pool size
                retryWrites=True
            )
            
            # Test the connection
            self.client.admin.command('ping')
            self.db = self.client[self.database_name]
            
            logger.info(f"Successfully connected to MongoDB database: {self.database_name}")
            return True
            
        except (ConnectionFailure, ServerSelectionTimeoutError) as e:
            logger.error(f"Failed to connect to MongoDB: {e}")
            return False
        except Exception as e:
            logger.error(f"Unexpected error connecting to MongoDB: {e}")
            return False
    
    def get_database(self):
        """Get database instance"""
        if self.db is None:
            if not self.connect():
                raise ConnectionError("Could not establish database connection")
        return self.db
    
    def get_collection(self, collection_name: str):
        """Get a specific collection"""
        db = self.get_database()
        return db[collection_name]
    
    def close(self):
        """Close database connection"""
        if self.client:
            self.client.close()
            logger.info("Database connection closed")

# Global database connection instance
db_connection = DatabaseConnection()

def get_db():
    """Get database instance - convenience function"""
    return db_connection.get_database()

def get_users_collection():
    """Get users collection"""
    return db_connection.get_collection('users')

def get_chat_sessions_collection():
    """Get chat sessions collection"""
    return db_connection.get_collection('chat_sessions')

def init_database():
    """Initialize database connection and create indexes"""
    try:
        db = get_db()
        
        if db is None:
            logger.error("Database connection is None")
            return False
        
        # Create indexes for better performance
        users_collection = get_users_collection()
        chat_sessions_collection = get_chat_sessions_collection()
        
        # Create unique index on email for users collection
        users_collection.create_index("email", unique=True)
        logger.info("Created unique index on users.email")
        
        # Create indexes on chat_sessions collection
        chat_sessions_collection.create_index("user_id")
        chat_sessions_collection.create_index("timestamp")
        chat_sessions_collection.create_index([("user_id", 1), ("timestamp", -1)])
        logger.info("Created indexes on chat_sessions collection")
        
        return True
        
    except Exception as e:
        logger.error(f"Failed to initialize database: {e}")
        return False

def test_connection():
    """Test database connection"""
    try:
        db = get_db()
        # Simple test query
        result = db.command('ping')
        logger.info("Database connection test successful")
        return True
    except Exception as e:
        logger.error(f"Database connection test failed: {e}")
        return False

if __name__ == "__main__":
    # Test the connection when run directly
    print("Testing MongoDB connection...")
    if test_connection():
        print("✓ Database connection successful")
        if init_database():
            print("✓ Database initialization successful")
        else:
            print("✗ Database initialization failed")
    else:
        print("✗ Database connection failed")