src/db.py

import os
import re
import time
from typing import Dict, List, Optional

try:
    import pymongo
    HAS_PYMONGO = True
except ImportError:
    HAS_PYMONGO = False

# Read connection string strictly from Hugging Face Space Secrets
MONGODB_URI = os.environ.get("MONGODB_URI", "").strip()
DB_NAME = os.environ.get("MONGODB_DB_NAME", "llm_xray_db")
COLLECTION_NAME = "audit_certificates"

_mongo_client = None


def get_db_collection(retries: int = 3, delay: float = 2.0):
    """Initializes and returns the MongoDB collection using HF Secrets.
    Retries the initial connection — HF Space cold starts often need a
    few extra seconds for DNS/TLS to Atlas before serverSelection succeeds.
    """
    global _mongo_client
    if not HAS_PYMONGO:
        return None
    if not MONGODB_URI:
        return None

    for attempt in range(1, retries + 1):
        try:
            if _mongo_client is None:
                _mongo_client = pymongo.MongoClient(
                    MONGODB_URI,
                    serverSelectionTimeoutMS=8000,
                    connectTimeoutMS=8000,
                )
            # Force a real round-trip here (pymongo is lazy otherwise),
            # so a slow cold-start connection is retried at this single
            # choke point rather than surfacing as "0 certificates".
            _mongo_client.admin.command("ping")
            return _mongo_client[DB_NAME][COLLECTION_NAME]
        except Exception as e:
            print(f"[MongoDB] Connection attempt {attempt}/{retries} warning: {e}", flush=True)
            _mongo_client = None
            if attempt < retries:
                time.sleep(delay)
    return None


def mongo_save_certificate(cert: Dict) -> bool:
    """Persists an audit certificate permanently to MongoDB Atlas."""
    col = get_db_collection()
    if col is None:
        return False
    try:
        config = cert.get("config", {})
        model_name = config.get("model_name", "unknown").strip()
        model_sha = str(config.get("model_sha", "main")).strip()
        doc_id = f"{model_name}__{model_sha}".replace("/", "__").lower()

        doc = dict(cert)
        doc["_id"] = doc_id
        doc["model_name"] = model_name
        doc["model_sha"] = model_sha
        doc["updated_at"] = cert.get("audited_at", "")

        col.replace_one({"_id": doc_id}, doc, upsert=True)
        print(f"[MongoDB] ✅ Successfully persisted '{model_name}' to Atlas collection", flush=True)
        return True
    except Exception as e:
        print(f"[MongoDB] ⚠️ Error saving certificate: {e}", flush=True)
        return False


def mongo_get_certificate(clean_model_name: str) -> Optional[Dict]:
    """Retrieves an audit certificate by model name from MongoDB."""
    col = get_db_collection()
    if col is None or not clean_model_name:
        return None
    try:
        norm = clean_model_name.strip()
        doc_id = norm.replace("/", "__").lower()
        doc = col.find_one({
            "$or": [
                {"_id": {"$regex": f"^{doc_id}", "$options": "i"}},
                {"config.model_name": {"$regex": f"^{re.escape(norm)}$", "$options": "i"}},
                {"model_name": {"$regex": f"^{re.escape(norm)}$", "$options": "i"}}
            ]
        })
        if doc:
            doc.pop("_id", None)
            return doc
    except Exception as e:
        print(f"[MongoDB] ⚠️ Error retrieving '{clean_model_name}': {e}", flush=True)
    return None


def mongo_get_all_certificates() -> List[Dict]:
    """Recalls all saved certificates."""
    col = get_db_collection()
    if col is None:
        return []
    try:
        docs = list(col.find({}))
        for d in docs:
            d.pop("_id", None)
        return docs
    except Exception as e:
        print(f"[MongoDB] ⚠️ Error loading certificates: {e}", flush=True)
        return []