"""
MongoDB Connection Helper
Generated by vai v{{vaiVersion}} on {{generatedAt}}

Database: {{db}}
Collection: {{collection}}
"""

import os
from datetime import datetime
from typing import List, Dict, Any, Optional
from pymongo import MongoClient

MONGODB_URI = os.getenv("MONGODB_URI")

if not MONGODB_URI:
    raise ValueError("MONGODB_URI environment variable is required")

# Module-level connection (singleton pattern)
_client: Optional[MongoClient] = None
_db = None


def connect():
    """Connect to MongoDB and return the database instance."""
    global _client, _db
    
    if _db is not None:
        return _db
    
    _client = MongoClient(MONGODB_URI)
    _db = _client["{{db}}"]
    print(f"Connected to MongoDB: {{db}}")
    return _db


def get_collection():
    """Get the documents collection."""
    db = connect()
    return db["{{collection}}"]


def close():
    """Close the MongoDB connection."""
    global _client, _db
    if _client:
        _client.close()
        _client = None
        _db = None


def vector_search(
    embedding: List[float],
    limit: int = 10,
    num_candidates: Optional[int] = None,
    filter_query: Optional[Dict] = None,
) -> List[Dict[str, Any]]:
    """
    Perform a vector search on the collection.
    
    Args:
        embedding: Query embedding vector
        limit: Number of results to return
        num_candidates: Candidates to consider (default: limit * 10)
        filter_query: Optional MongoDB filter
        
    Returns:
        List of documents with scores
    """
    collection = get_collection()
    num_candidates = num_candidates or limit * 10
    
    pipeline = [
        {
            "$vectorSearch": {
                "index": "{{index}}",
                "path": "{{field}}",
                "queryVector": embedding,
                "numCandidates": num_candidates,
                "limit": limit,
            }
        },
        {
            "$project": {
                "_id": 1,
                "text": 1,
                "metadata": 1,
                "score": {"$meta": "vectorSearchScore"},
            }
        },
    ]
    
    # Add filter if provided
    if filter_query:
        pipeline[0]["$vectorSearch"]["filter"] = filter_query
    
    results = list(collection.aggregate(pipeline))
    return [{"document": doc, "score": doc.pop("score")} for doc in results]


def insert_documents(docs: List[Dict[str, Any]]) -> Dict[str, int]:
    """
    Insert documents with embeddings.
    
    Args:
        docs: List of dicts with 'text', 'embedding', and optional 'metadata'
        
    Returns:
        Dict with 'inserted_count'
    """
    collection = get_collection()
    
    documents = [
        {
            "text": doc["text"],
            "{{field}}": doc["embedding"],
            "metadata": doc.get("metadata", {}),
            "_embeddedAt": datetime.utcnow(),
            "_model": "{{model}}",
        }
        for doc in docs
    ]
    
    result = collection.insert_many(documents)
    return {"inserted_count": len(result.inserted_ids)}