from typing import List, Optional, Union, Any
import dspy
import os
from openai import (
    OpenAI,
    APITimeoutError,
    InternalServerError,
    RateLimitError,
    UnprocessableEntityError,
)
import backoff

try:
    from pymongo import MongoClient
    from pymongo.errors import (
        ConnectionFailure,
        ConfigurationError,
        ServerSelectionTimeoutError,
        InvalidURI,
        OperationFailure,
    )
except ImportError:
    raise ImportError(
        "Please install the pymongo package by running `pip install dspy-ai[mongodb]`"
    )


def build_vector_search_pipeline(
    index_name: str, query_vector: List[float], num_candidates: int, limit: int
) -> List[dict[str, Any]]:
    return [
        {
            "$vectorSearch": {
                "index": index_name,
                "path": "embedding",
                "queryVector": query_vector,
                "numCandidates": num_candidates,
                "limit": limit,
            }
        },
        {"$project": {"_id": 0, "text": 1, "score": {"$meta": "vectorSearchScore"}}},
    ]


class Embedder:
    def __init__(self, provider: str, model: str):
        if provider == "openai":
            api_key = os.getenv("OPENAI_API_KEY")
            if not api_key:
                raise ValueError("Environment variable OPENAI_API_KEY must be set")
            self.client = OpenAI()
            self.model = model

    @backoff.on_exception(
        backoff.expo,
        (
            APITimeoutError,
            InternalServerError,
            RateLimitError,
            UnprocessableEntityError,
        ),
        max_time=15,
    )
    def __call__(self, queries) -> Any:
        embedding = self.client.embeddings.create(input=queries, model=self.model)
        return [result.embedding for result in embedding.data]


class MongoDBAtlasRM(dspy.Retrieve):
    def __init__(
        self,
        db_name: str,
        collection_name: str,
        index_name: str,
        k: int = 5,
        embedding_provider: str = "openai",
        embedding_model: str = "text-embedding-ada-002",
    ):
        super().__init__(k=k)
        self.db_name = db_name
        self.collection_name = collection_name
        self.index_name = index_name
        self.username = os.getenv("ATLAS_USERNAME")
        self.password = os.getenv("ATLAS_PASSWORD")
        self.cluster_url = os.getenv("ATLAS_CLUSTER_URL")
        if not self.username:
            raise ValueError("Environment variable ATLAS_USERNAME must be set")
        if not self.password:
            raise ValueError("Environment variable ATLAS_PASSWORD must be set")
        if not self.cluster_url:
            raise ValueError("Environment variable ATLAS_CLUSTER_URL must be set")
        try:
            self.client = MongoClient(
                f"mongodb+srv://{self.username}:{self.password}@{self.cluster_url}/{self.db_name}"
                "?retryWrites=true&w=majority"
            )
        except (
            InvalidURI,
            ConfigurationError,
            ConnectionFailure,
            ServerSelectionTimeoutError,
            OperationFailure,
        ) as e:
            raise ConnectionError("Failed to connect to MongoDB Atlas") from e

        self.embedder = Embedder(provider=embedding_provider, model=embedding_model)

    def forward(self, query_or_queries: str) -> dspy.Prediction:
        query_vector = self.embedder([query_or_queries])
        pipeline = build_vector_search_pipeline(
            index_name=self.index_name,
            query_vector=query_vector[0],
            num_candidates=self.k * 10,
            limit=self.k,
        )
        contents = self.client[self.db_name][self.collection_name].aggregate(pipeline)
        return dspy.Prediction(passages=list(contents))