src/Database/operations.py

from __future__ import annotations

from datetime import datetime
from typing import Any, Dict, List, Optional

from pymongo import DESCENDING
from Database.connections import get_img_gen_collection
import streamlit as st
from core.logger import logger


def insert_img_gen_run_result(
        *,
        prompts: List[str],
        image_urls: List[str],
        category: str = "general",
        platform: Optional[str] = None,
        model_name: Optional[str] = None,
        offer: Optional[str] = None,
        framework: Optional[str] = None,
        user_query: Optional[str] = None,
        created_by: Optional[str] = None,
) -> str:
    col = get_img_gen_collection()
    if col is None:
        raise ValueError("img_gen collection not available.")

    if not prompts or not image_urls:
        raise ValueError("prompts and image_urls cannot be empty")

    doc: Dict[str, Any] = {
        "source": "img_generator",
        "category": category or "general",
        "prompts": prompts,
        "image_urls": image_urls,
        "platform": platform,
        "model_name": model_name,
        "offer": offer if offer else '',
        "framework": framework if framework else '',
        "user_query": user_query,
        "created_by": created_by,
        "created_at": datetime.utcnow(),
    }
    res = col.insert_one(doc)
    return str(res.inserted_id)


def find_img_gen_results(
        *,
        category: Optional[str] = None,
        start_date: Optional[datetime] = None,
        end_date: Optional[datetime] = None,
        limit: int = 200,
        created_by: Optional[str] = None,
) -> List[Dict[str, Any]]:
    try:
        col = get_img_gen_collection()
        if col is None:
            return []

        q: Dict[str, Any] = {"source": "img_generator"}

        if category:
            q["category"] = category

        if start_date or end_date:
            rng: Dict[str, Any] = {}
            if start_date:
                rng["$gte"] = start_date
            if end_date:
                rng["$lt"] = end_date
            q["created_at"] = rng

        if created_by:
            q["created_by"] = created_by

        cur = (
            col.find(q)
            .sort("created_at", DESCENDING)
            .limit(limit)
        )

        out: List[Dict[str, Any]] = []
        for d in cur:
            d["_id"] = str(d.get("_id"))
            out.append(d)

        return out
    except Exception as e:
        st.error(f"Operations File: {e}")


def list_img_gen_categories(*, created_by: Optional[str] = None) -> List[str]:
    col = get_img_gen_collection()
    if col is None:
        return []

    q: Dict[str, Any] = {"source": "img_generator"}
    if created_by:
        q["created_by"] = created_by

    cats = col.distinct("category", q)
    cats = [c for c in cats if c]
    cats.sort()
    return cats