from motor.motor_asyncio import AsyncIOMotorClient
from pymongo import MongoClient
import os
import logging

# Global Motor client singleton - per event loop
_motor_clients = {}  # event_loop_id -> (client, database)

async def get_db():
    """Get database connection using singleton Motor client per event loop."""
    import asyncio
    
    # Get current event loop to ensure Motor client affinity
    try:
        current_loop = asyncio.get_running_loop()
        loop_id = id(current_loop)
    except RuntimeError:
        raise RuntimeError("get_db() must be called from within an async context")
    
    # Return cached database for this event loop if available
    if loop_id in _motor_clients:
        client, database = _motor_clients[loop_id]
        return database
    
    # Read MongoDB connection from environment
    mongo_uri = os.environ.get('MONGO_URI')
    mongo_user = os.environ.get('MONGO_USER')
    mongo_pass = os.environ.get('MONGO_PASS')
    mongo_host = os.environ.get('MONGO_HOST', 'localhost')
    mongo_port = os.environ.get('MONGO_PORT', '27017')

    # Build URI: prefer MONGO_URI, fall back to host+port with optional credentials
    if not mongo_uri:
        if mongo_user and mongo_pass:
            mongo_uri = f"mongodb://{mongo_user}:{mongo_pass}@{mongo_host}:{mongo_port}/semblance_db?authSource=admin"
        else:
            mongo_uri = f"mongodb://{mongo_host}:{mongo_port}"

    try:
        motor_client = AsyncIOMotorClient(mongo_uri, serverSelectionTimeoutMS=5000)
        database = motor_client.semblance_db
        await database.command('ping')
        logging.info("Successfully connected to MongoDB")
    except Exception as e:
        raise RuntimeError(f"Failed to connect to MongoDB: {e}. Check MONGO_URI in backend/.env.") from e

    _motor_clients[loop_id] = (motor_client, database)

    # Ensure indexes exist (idempotent)
    try:
        await database.users.create_index("username", unique=True, background=True)
        await database.users.create_index("email", unique=True, background=True)
        await database.users.create_index("role", background=True)
        await database.users.create_index([("is_active", 1), ("username", 1)], background=True)
        await database.personas.create_index("created_by", background=True)
        await database.focus_groups.create_index("created_by", background=True)
        await database.folders.create_index("created_by", background=True)
        await database.folders.create_index("parent_folder_id", background=True)
        # usage_events indexes
        await database.usage_events.create_index([("user_id", 1), ("ts", -1)], background=True)
        await database.usage_events.create_index([("focus_group_id", 1), ("ts", -1)], background=True)
        await database.usage_events.create_index([("ts", -1)], background=True)
        await database.usage_events.create_index([("feature", 1), ("ts", -1)], background=True)
        await database.usage_events.create_index([("model", 1), ("ts", -1)], background=True)
        await database.usage_events.create_index([("status", 1), ("ts", -1)], background=True)
        # model_pricing indexes
        await database.model_pricing.create_index([("model", 1), ("effective_from", -1)], background=True)
    except Exception as e:
        logging.warning(f"Index creation warning (non-fatal): {e}")

    return database


def close_db_connections():
    """Close all Motor clients and their PyMongo background threads."""
    global _motor_clients
    
    closed_count = 0
    for loop_id, (client, database) in _motor_clients.items():
        try:
            client.close()
            closed_count += 1
        except Exception as e:
            logging.warning(f"Error closing Motor client for loop {loop_id}: {e}")
    
    if closed_count > 0:
        logging.info(f"🗄️  Closed {closed_count} Motor clients - PyMongo threads should stop")
    
    _motor_clients.clear()