from pymongo import AsyncMongoClient


class AsyncMongo:
    def __init__(
        self, mongo_uri: str, database_name: str = "MayaContigo", mode: str = "cosmosdb"
    ) -> None:
        self.mode = mode
        self.client: AsyncMongoClient = AsyncMongoClient(mongo_uri)
        self.db = self.client.get_database(database_name)

    def build_pipeline(self, embedding, limit):
        if self.mode == "native":
            return {
                "$vectorSearch": {
                    "index": "vector_index",
                    "queryVector": embedding,
                    "path": "embedding",
                    "limit": limit,
                    "numCandidates": 3 * limit,
                }
            }
        elif self.mode == "cosmosdb":
            return {
                "$search": {
                    "cosmosSearch": {
                        "path": "vector",
                        "vector": embedding[:2000],
                        "k": limit,
                    }
                }
            }
        else:
            raise ValueError("Invalid mode")

    async def semantic_search(
        self,
        collection: str,
        embedding: list[float],
        limit: int = 10,
        conditions: dict | None = None,
        threshold: float | None = None,
        **kwargs,
    ):
        db_collection = self.db[collection]

        pipeline = self.build_pipeline(embedding, limit)
        aggregate = await db_collection.aggregate([pipeline])
        return [result async for result in aggregate]
