from motor.motor_asyncio import AsyncIOMotorClient
from app.config import settings

class MongoDB:
    client: AsyncIOMotorClient = None
    db = None

    def connect(self) -> None:
        """Establish MongoDB connection client."""
        self.client = AsyncIOMotorClient(settings.MONGODB_URI)
        self.db = self.client[settings.MONGODB_DB]

    def disconnect(self) -> None:
        """Close MongoDB connection client."""
        if self.client:
            self.client.close()

    @property
    def email_bodies(self):
        return self.db["email_bodies"]

    @property
    def agent_memory(self):
        return self.db["agent_memory"]

    @property
    def agent_logs(self):
        return self.db["agent_logs"]

# Global MongoDB wrapper
mongo_db = MongoDB()

async def init_mongo() -> None:
    """Initialize collections and indexes."""
    mongo_db.connect()
    # Create indexes for high performance querying
    await mongo_db.email_bodies.create_index("email_id", unique=True)
    await mongo_db.agent_memory.create_index([("user_id", 1), ("email_id", 1)])
    await mongo_db.agent_logs.create_index([("run_id", 1), ("step", 1)])