database/connections.py

from pymongo import MongoClient
from core.logger import logger
from core.settings import cnf

import streamlit as st

@st.cache_resource
def get_mongo_client():
    try:
        uri = cnf.MONGO_URI
        dbname = cnf.MONGO_DB
        
        # Check if required environment variables are set
        if not uri or not dbname:
            logger.warning("MongoDB configuration not found. Running without database.")
            return None
            
        client = MongoClient(uri)
        db = client[dbname]
        db.command("ping")
        logger.info("MongoDB connection established successfully")
        return db
    except Exception as e:
        logger.error(f"MongoDB connection failed: {e}")
        return None

def get_results_collection():
    db = get_mongo_client()
    if db is None:
        return None
    return db["img_gen"]

def get_image_collection():
    db = get_mongo_client()
    if db is None:
        return None
    return db["image_analyses"]


def get_video_collection():
    db = get_mongo_client()
    if db is None:
        return None
    return db["video_analyses"]

def get_script_collection():
    db = get_mongo_client()
    if db is None:
        return None
    return db["script_generation"]