from qdrant_client import AsyncQdrantClient, models
from openai import AsyncOpenAI


class LamehVectorManager:
    def __init__(
        self,
        qdrant_url: str,
        openai_api_key: str,
        collection_name: str,
        vector_size: int,
    ):
        """
        Initialize the Lameh Vector Manager with Qdrant and OpenAI configurations.
        """
        self.qdrant_client = AsyncQdrantClient(url=qdrant_url)
        self.openai_client = AsyncOpenAI(api_key=openai_api_key)
        self.collection_name = collection_name
        self.vector_size = vector_size

    async def initialize_collection(self, distance_metric: str = "COSINE"):
        """
        Create a collection in Qdrant if it doesn't exist.
        """
        try:
            await self.qdrant_client.create_collection(
                collection_name=self.collection_name,
                vectors_config=models.VectorParams(
                    size=self.vector_size, distance=models.Distance[distance_metric]
                ),
            )
        except Exception as e:
            print(f"Collection might already exist or an error occurred: {e}")

    async def upsert_vectors(self, items: list[dict]):
        """
        Upserts vectors into the Qdrant collection.
        :param items: A list of dictionaries containing 'id', 'data', and optional 'payload'.
        """
        try:
            points = [
                models.PointStruct(
                    id=item["id"],
                    vector=(await self.generate_embedding(item["data"]))
                    .data[0]
                    .embedding,
                    payload=item.get("payload", {}),
                )
                for item in items
            ]
            await self.qdrant_client.upsert(
                collection_name=self.collection_name, points=points
            )
        except Exception as e:
            print(f"Error during upsert: {e}")

    async def search(self, query: str, limit: int = 10, threshold: float = 0.9):
        """
        Searches for the most similar vectors to the given query.
        :param query: The search query.
        :param limit: Number of results to retrieve.
        """
        try:
            query_vector = (await self.generate_embedding(query)).data[0].embedding
            return await self.qdrant_client.search(
                collection_name=self.collection_name,
                query_vector=query_vector,
                limit=limit,
                with_payload=True,
                with_vectors=True,
                score_threshold=threshold,
            )
        except Exception as e:
            print(f"Error during search: {e}")
            return []

    async def close(self):
        """
        Close the Qdrant connection.
        """
        await self.qdrant_client.close()

    async def generate_embedding(self, data: str):
        """
        Generate embeddings using OpenAI's embedding API.
        """
        return await self.openai_client.embeddings.create(
            input=[data], model="text-embedding-3-large"
        )
