modules/database/database_init.py

import os
import logging
from azure.cosmos import CosmosClient
from pymongo import MongoClient
import certifi

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

# Variables globales para Cosmos DB SQL API
cosmos_client = None
user_database = None
user_container = None
application_requests_container = None
user_feedback_container = None
user_sessions_container = None  # Agregado para sessions

# Variables globales para Cosmos DB MongoDB API
mongo_client = None
mongo_db = None

def verify_container_partition_key(container, expected_path):
    """Verifica la configuración de partition key de un contenedor"""
    try:
        container_props = container.read()
        partition_key_paths = container_props['partitionKey']['paths']
        
        logger.info(f"Container: {container.id}, Partition Key Paths: {partition_key_paths}")
        
        if expected_path not in partition_key_paths:
            logger.warning(
                f"Partition key incorrecta en {container.id}. "
                f"Actual: {partition_key_paths}, Esperada: {expected_path}"
            )
            return False
        return True
    except Exception as e:
        logger.error(f"Error verificando partition key en {container.id}: {str(e)}")
        return False

def initialize_cosmos_sql_connection():
    """Inicializa y verifica la conexión a Cosmos DB SQL API"""
    global cosmos_client, user_database, user_container, \
           application_requests_container, user_feedback_container, \
           user_sessions_container

    try:
        cosmos_endpoint = os.environ.get("COSMOS_ENDPOINT")
        cosmos_key = os.environ.get("COSMOS_KEY")
        
        if not cosmos_endpoint or not cosmos_key:
            raise ValueError("COSMOS_ENDPOINT and COSMOS_KEY environment variables must be set")

        # Inicializar cliente y obtener contenedores
        cosmos_client = CosmosClient(cosmos_endpoint, cosmos_key)
        user_database = cosmos_client.get_database_client("user_database")
        
        # Configuración de contenedores y sus partition keys esperadas
        containers_config = {
            "users": ("/id", user_container),
            "application_requests": ("/email", application_requests_container),
            "user_feedback": ("/username", user_feedback_container),
            "users_sessions": ("/username", user_sessions_container)
        }

        # Inicializar y verificar cada contenedor
        for container_id, (expected_path, container_ref) in containers_config.items():
            container = user_database.get_container_client(container_id)
            if not verify_container_partition_key(container, expected_path):
                logger.error(f"Verificación de partition key falló para {container_id}")
                return False
                
            # Actualizar la referencia global
            if container_id == "users":
                user_container = container
            elif container_id == "application_requests":
                application_requests_container = container
            elif container_id == "user_feedback":
                user_feedback_container = container
            elif container_id == "users_sessions":
                user_sessions_container = container

        logger.info("Conexión a Cosmos DB SQL API exitosa y particiones verificadas")
        return True

    except Exception as e:
        logger.error(f"Error al conectar con Cosmos DB SQL API: {str(e)}", exc_info=True)
        return False

def initialize_mongodb_connection():
    """Inicializa la conexión a Cosmos DB MongoDB API"""
    global mongo_client, mongo_db
    try:
        cosmos_mongodb_connection_string = os.getenv("MONGODB_CONNECTION_STRING")
        if not cosmos_mongodb_connection_string:
            raise ValueError("MONGODB_CONNECTION_STRING environment variable is not set")

        mongo_client = MongoClient(
            cosmos_mongodb_connection_string,
            tls=True,
            tlsCAFile=certifi.where(),
            retryWrites=False,
            serverSelectionTimeoutMS=5000,
            connectTimeoutMS=10000,
            socketTimeoutMS=10000
        )
        
        # Verificar conexión
        mongo_client.admin.command('ping')
        mongo_db = mongo_client['aideatext_db']
        
        logger.info("Conexión a Cosmos DB MongoDB API exitosa")
        return True
        
    except Exception as e:
        logger.error(f"Error al conectar con Cosmos DB MongoDB API: {str(e)}", exc_info=True)
        return False

def initialize_database_connections():
    """Inicializa todas las conexiones a bases de datos"""
    sql_success = initialize_cosmos_sql_connection()
    mongodb_success = initialize_mongodb_connection()
    
    if sql_success and mongodb_success:
        logger.info("Todas las conexiones a bases de datos inicializadas correctamente")
    else:
        logger.error("Falló la inicialización de una o más conexiones a bases de datos")
        
    return sql_success and mongodb_success

def get_sql_containers():
    """Obtiene los contenedores SQL, reinicializando si es necesario"""
    if any(container is None for container in 
           [user_container, application_requests_container, 
            user_feedback_container, user_sessions_container]):
        initialize_cosmos_sql_connection()
    return user_container, application_requests_container, user_feedback_container, user_sessions_container

def get_mongodb():
    """Obtiene la conexión MongoDB, reinicializando si es necesario"""
    if mongo_db is None:
        initialize_mongodb_connection()
    return mongo_db