core/db.py

from motor.motor_asyncio import AsyncIOMotorClient
from qdrant_client import AsyncQdrantClient
from .config import settings

class Database:
    mongo_client: AsyncIOMotorClient = None
    qdrant_client: AsyncQdrantClient = None
    db = None

    async def connect(self):
        # Connect to MongoDB
        if settings.MONGO_URI:
            self.mongo_client = AsyncIOMotorClient(settings.MONGO_URI)
            self.db = self.mongo_client.janrakshak
            
            # Optimization: Create a geospatial index for efficient hotspot queries
            import pymongo
            try:
                # serverSelectionTimeoutMS is set to 5000 so it doesn't hang the server startup for 30 seconds
                self.mongo_client = AsyncIOMotorClient(settings.MONGO_URI, serverSelectionTimeoutMS=5000)
                self.db = self.mongo_client.janrakshak
                await self.db["incidents"].create_index([("location", pymongo.GEOSPHERE)])
                print("Connected to MongoDB and verified indexes")
            except Exception as e:
                print(f"Warning: Could not connect to MongoDB. Is the MONGO_URI correct? Error: {e}")
        
        # Connect to Qdrant
        if settings.QDRANT_URL:
            self.qdrant_client = AsyncQdrantClient(
                url=settings.QDRANT_URL,
                api_key=settings.QDRANT_API_KEY if settings.QDRANT_API_KEY else None
            )
            print("Connected to Qdrant")

    async def disconnect(self):
        if self.mongo_client:
            self.mongo_client.close()
            print("Disconnected from MongoDB")
        if self.qdrant_client:
            await self.qdrant_client.close()
            print("Disconnected from Qdrant")

    import hashlib
    import json
    
    async def generate_evidence_hash(self, payload: dict, file_path: str = None) -> str:
        """
        Creates a legally admissible SHA-256 hash of the evidence metadata 
        and the raw file bytes to prove it hasn't been tampered with.
        """
        hasher = self.hashlib.sha256()
        
        # Hash the metadata securely
        payload_str = self.json.dumps(payload, sort_keys=True, default=str)
        hasher.update(payload_str.encode('utf-8'))
        
        # If there's an associated file (audio/image), hash its raw bytes too
        if file_path:
            import aiofiles
            try:
                async with aiofiles.open(file_path, 'rb') as f:
                    while chunk := await f.read(8192):
                        hasher.update(chunk)
            except Exception as e:
                print(f"Warning: Could not hash file at {file_path}: {e}")
                
        return hasher.hexdigest()

db_handler = Database()