database.py

import os
from datetime import datetime
from pymongo import MongoClient
from dotenv import load_dotenv
from bson import ObjectId

# Load environment variables
load_dotenv()

# MongoDB Configuration
MONGO_URI = os.getenv("MONGO_URI")
DB_NAME = os.getenv("MONGO_DB")

# Initialize MongoDB client
client = MongoClient(MONGO_URI)
db = client[DB_NAME]


video_collection = db["video_analysis"]
script_collection = db["script_generator"]
comparison_collection = db["comparison_results"]


#  VIDEO ANALYSIS 
def insert_video_analysis(video_name: str, response: dict,thumbnail:str):
    """Insert a new video analysis result"""
    document = {
        "video_name": video_name,
        "response": response,
        "thumbnail": thumbnail,
        "created_at": datetime.now(),
        "lob": "test"
    }
    result = video_collection.insert_one(document)
    return str(result.inserted_id)


def get_all_video_analyses(limit: int = 20):
    """Fetch all saved video analyses"""
    return list(video_collection.find().sort("created_at", -1).limit(limit))

#SCRIPT GENERATOR
def insert_script_result(
    video_name: str,
    offer_details: str,
    target_audience: str,
    specific_hook: str,
    additional_context: str,
    response: dict,
    thumbnail:str
):
    """Insert a new script generation result"""
    document = {
        "video_name": video_name,
        "offer_details": offer_details,
        "target_audience": target_audience,
        "specific_hook": specific_hook,
        "additional_context": additional_context,
        "response": response,
        "thumbnail":thumbnail,
        "type": "script",
        "created_at": datetime.now(),
        "lob": "test"
    }
    result = script_collection.insert_one(document)
    return str(result.inserted_id)


def get_all_scripts(limit: int = 20):
    """Fetch all saved script generations"""
    return list(script_collection.find({"type": "script"}).sort("created_at", -1).limit(limit))



# COMPARISON 
def insert_comparison_result(
    video_names: list,
    user_prompt: str,
    response: dict,
    video_name: str = "comparison_result",
    thumbnails: dict = None
):
    """Insert comparison result between videos"""
    document = {
        "video_name": video_name,
        "video_names": video_names,
        "user_prompt": user_prompt,
        "response": response,
        "type": "comparison", 
        "thumbnails": thumbnails or {},   
        "created_at": datetime.now(),
        "lob": "test"
    }
    result = comparison_collection.insert_one(document)
    return str(result.inserted_id)



def get_all_comparisons(limit: int = 20):
    """Fetch all saved comparison results"""
    return list(comparison_collection.find().sort("created_at", -1).limit(limit))



def get_result_by_id(collection, doc_id):
    """Fetch a specific result by _id from any collection"""
    return collection.find_one({"_id": ObjectId(doc_id)})


def delete_result_by_id(collection, doc_id):
    """Delete a specific result by _id from any collection"""
    result = collection.delete_one({"_id": ObjectId(doc_id)})
    return result.deleted_count > 0