mongodb_logging.py

import logging
import os
from datetime import datetime
from typing import Dict, Any
from pymongo import MongoClient
from pymongo.collection import Collection


class MongoDBLogHandler(logging.Handler):
    """Custom logging handler that stores logs in MongoDB"""
    
    def __init__(self, collection: Collection):
        super().__init__()
        self.collection = collection
        
    def emit(self, record: logging.LogRecord):
        """Emit a log record to MongoDB"""
        try:
            # Create log document
            log_doc = {
                "timestamp": datetime.utcnow(),
                "level": record.levelname,
                "logger": record.name,
                "message": record.getMessage(),
                "module": record.module,
                "function": record.funcName,
                "line": record.lineno,
                "thread": record.thread,
                "process": record.process,
            }
            
            # Add exception info if present
            if record.exc_info:
                log_doc["exception"] = self.formatException(record.exc_info)
                
            # Add extra fields if present
            if hasattr(record, 'extra_fields'):
                log_doc.update(record.extra_fields)
                
            # Insert into MongoDB
            self.collection.insert_one(log_doc)
            
        except Exception as e:
            # Fallback to console if MongoDB fails
            print(f"MongoDB logging failed: {e}")
            print(f"Original log: {record.getMessage()}")


def setup_mongodb_logging(mongo_uri: str, database_name: str = "HairSwapDB", collection_name: str = "logs"):
    """Setup MongoDB logging for the application"""
    
    try:
        # Connect to MongoDB
        client = MongoClient(mongo_uri)
        db = client.get_database(database_name)
        logs_collection = db.get_collection(collection_name)
        
        # Create MongoDB log handler
        mongo_handler = MongoDBLogHandler(logs_collection)
        mongo_handler.setLevel(logging.INFO)
        
        # Create formatter
        formatter = logging.Formatter(
            '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        )
        mongo_handler.setFormatter(formatter)
        
        # Get root logger and add handler
        root_logger = logging.getLogger()
        root_logger.addHandler(mongo_handler)
        
        # Also add to specific logger
        hair_logger = logging.getLogger("hair_server")
        hair_logger.addHandler(mongo_handler)
        
        print(f"✅ MongoDB logging enabled - storing logs in {database_name}.{collection_name}")
        return True
        
    except Exception as e:
        print(f"❌ Failed to setup MongoDB logging: {e}")
        return False


def get_logs_from_mongodb(mongo_uri: str, database_name: str = "HairSwapDB", 
                         collection_name: str = "logs", limit: int = 100, 
                         level: str = None, logger_name: str = None):
    """Retrieve logs from MongoDB"""
    
    try:
        client = MongoClient(mongo_uri)
        db = client.get_database(database_name)
        logs_collection = db.get_collection(collection_name)
        
        # Build query
        query = {}
        if level:
            query["level"] = level
        if logger_name:
            query["logger"] = logger_name
            
        # Get logs sorted by timestamp (newest first)
        logs = list(logs_collection.find(query)
                   .sort("timestamp", -1)
                   .limit(limit))
        
        # Convert ObjectId to string for JSON serialization
        for log in logs:
            log["_id"] = str(log["_id"])
            if "timestamp" in log:
                log["timestamp"] = log["timestamp"].isoformat()
                
        return logs
        
    except Exception as e:
        print(f"❌ Failed to retrieve logs from MongoDB: {e}")
        return []


def clear_logs_from_mongodb(mongo_uri: str, database_name: str = "HairSwapDB", 
                          collection_name: str = "logs", days_older_than: int = None):
    """Clear old logs from MongoDB"""
    
    try:
        client = MongoClient(mongo_uri)
        db = client.get_database(database_name)
        logs_collection = db.get_collection(collection_name)
        
        if days_older_than:
            from datetime import timedelta
            cutoff_date = datetime.utcnow() - timedelta(days=days_older_than)
            result = logs_collection.delete_many({"timestamp": {"$lt": cutoff_date}})
            print(f"✅ Deleted {result.deleted_count} logs older than {days_older_than} days")
        else:
            result = logs_collection.delete_many({})
            print(f"✅ Deleted all {result.deleted_count} logs")
            
        return result.deleted_count
        
    except Exception as e:
        print(f"❌ Failed to clear logs from MongoDB: {e}")
        return 0