import os
import time
from typing import List, Dict, Optional
from pymongo import MongoClient, errors
from pymongo.collection import Collection

class MongoDBHandler:
    """
    Handles interactions with MongoDB for the Gemma-3 RAG system.
    Follows PythonMind principles: Simple, Modular, and Robust.
    """
    def __init__(self, uri: str = "mongodb://localhost:27017/", db_name: str = "rag_db", collection_name: str = "chunks"):
        self.uri = uri
        self.db_name = db_name
        self.collection_name = collection_name
        self.client: Optional[MongoClient] = None
        self.db = None
        self.collection: Optional[Collection] = None

    def connect(self) -> bool:
        """Establishes connection to MongoDB."""
        try:
            self.client = MongoClient(self.uri, serverSelectionTimeoutMS=5000)
            # Trigger a server selection to verify connection
            self.client.server_info()
            self.db = self.client[self.db_name]
            self.collection = self.db[self.collection_name]
            print(f"[MONGO] Connected to {self.uri}, Database: {self.db_name}, Collection: {self.collection_name}")
            return True
        except errors.ServerSelectionTimeoutError as e:
            print(f"[MONGO ERROR] Could not connect to MongoDB: {e}")
            return False
        except Exception as e:
            print(f"[MONGO ERROR] An unexpected error occurred: {e}")
            return False

    def insert_chunk(self, content: str, metadata: Dict) -> bool:
        """Inserts a single chunk into the collection."""
        if self.collection is None:
            if not self.connect(): return False
        
        try:
            document = {
                "content": content,
                "metadata": metadata,
                "timestamp": time.time()
            }
            self.collection.insert_one(document)
            return True
        except Exception as e:
            print(f"[MONGO ERROR] Insertion failed: {e}")
            return False

    def find_relevant(self, query_dict: Dict, limit: int = 5) -> List[Dict]:
        """Performs a standard query search."""
        if self.collection is None:
            if not self.connect(): return []
            
        try:
            results = self.collection.find(query_dict).limit(limit)
            return list(results)
        except Exception as e:
            print(f"[MONGO ERROR] Search failed: {e}")
            return []

    def clear(self) -> bool:
        """Clears the collection."""
        if self.collection is None:
            if not self.connect(): return False
        try:
            self.collection.delete_many({})
            return True
        except Exception as e:
            print(f"[MONGO ERROR] Clear failed: {e}")
            return False

    def get_stats(self) -> Dict:
        """Returns collection statistics."""
        if self.collection is None:
            if not self.connect(): return {"error": "Not connected"}
        try:
            count = self.collection.count_documents({})
            return {"count": count, "collection": self.collection_name, "db": self.db_name}
        except Exception as e:
            return {"error": str(e)}

    def close(self):
        """Closes the connection."""
        if self.client:
            self.client.close()