"""Qdrant-backed vector database implementation.

This module provides ``QdrantVectorDatabase``, a vector database that
delegates storage and similarity search to a Qdrant server (either
remote or local on-disk). Informations are serialized into Qdrant
point payloads via Base64 encoding, and cosine similarity is used for
nearest-neighbour search. The ``qdrant-client`` package must be
installed to use this backend.
"""

import os
import uuid
from typing import Callable, List, Optional, Union

import numpy as np
from arbol import aprint
from pydantic import BaseModel

from litemind.agent.augmentations.information.information import Information
from litemind.agent.augmentations.vector_db.default_vector_db import (
    DefaultVectorDatabase,
)
from litemind.apis.base_api import BaseApi
from litemind.media.media_base import MediaBase


class QdrantVectorDatabase(DefaultVectorDatabase):
    """
    A vector database implementation using Qdrant.

    This implementation uses the Qdrant vector database to store and retrieve informations.
    It requires the qdrant-client package to be installed.
    """

    def __init__(
        self,
        name: str = "QdrantVectorDB",
        description: Optional[str] = None,
        collection_name: str = "informations",
        embedding_function: Optional[
            Callable[[List[Information]], List[List[float]]]
        ] = None,
        embedding_length: Optional[int] = None,
        url: Optional[str] = None,
        location: Optional[str] = None,
        qdrant_api_key: Optional[str] = None,
        api: BaseApi = None,
    ):
        """
        Initialize the Qdrant vector database.

        Parameters
        ----------
        name: str
            The name of the augmentation.
        description: Optional[str]
            A description of the augmentation.
        collection_name: str
            The name of the collection in Qdrant.
        embedding_function: Optional[Callable[[List[Information]], List[List[float]]]]
            A function that takes a list of informations and returns a list of embedding vectors.
            If not provided, a default embedding function will be used.
        embedding_length: int
            The length of the embedding vectors.
        url: Optional[str]
            The URL of the Qdrant server. If None, will use a local instance.
        location: Optional[str]
            The location of the local Qdrant instance. Only used if url is None.
        qdrant_api_key: Optional[str]
            The API key for the Qdrant server.
        api: BaseApi
            The API to use for embedding. If not provided, uses litemind's default API.

        """
        super().__init__(
            name=name,
            description=description,
            embedding_function=embedding_function,
            embedding_length=embedding_length,
            api=api,
        )

        self.collection_name = collection_name

        try:
            from qdrant_client import QdrantClient
            from qdrant_client.http import models

            if url:
                self.client = QdrantClient(url=url, api_key=qdrant_api_key)
            else:
                if location is None:
                    # Get the user's home folder:
                    home_folder = os.path.expanduser("~")

                    # Qdrant data folder:
                    location = os.path.join(home_folder, ".qdrant_data")

                    # Create the Qdrant data folder:
                    os.makedirs(location, exist_ok=True)

                # Use the local Qdrant instance at the specified location:
                self.client = QdrantClient(path=location)

            # Create collection if it doesn't exist
            collections = self.client.get_collections().collections
            collection_names = [collection.name for collection in collections]

            if collection_name not in collection_names:
                aprint(f"Creating new Qdrant collection: {collection_name}")
                self.client.create_collection(
                    collection_name=collection_name,
                    vectors_config=models.VectorParams(
                        size=self.embedding_length,
                        distance=models.Distance.COSINE,
                    ),
                )

        except ImportError:
            raise ImportError(
                "Qdrant not installed. Please install it with 'pip install qdrant-client'."
            )

    def add_informations(
        self,
        informations: List[Information],
        embeddings: Optional[List[List[float]]] = None,
    ) -> List[str]:
        """
        Add informations to the database.

        Parameters
        ----------
        informations: List[Information]
            The informations to add to the database.
        embeddings: Optional[List[List[float]]]
            The embeddings for the informations. If not provided, they will be computed.

        Returns
        -------
        List[str]
            The IDs of the added informations.
        """
        from qdrant_client.http import models

        # Generate IDs for informations without IDs
        for info in informations:
            if info.id is None:
                info.set_id(str(uuid.uuid4()))

        # Create embeddings if not provided
        if embeddings is None:
            aprint(f"Computing embeddings for {len(informations)} informations")
            embeddings = self.embedding_function(informations)

        # Prepare points for Qdrant
        points = []
        for info, emb in zip(informations, embeddings):
            # Create a point
            payload = {"info": info.to_base64()}
            points.append(
                models.PointStruct(
                    id=info.id,
                    vector=emb,
                    payload=payload,
                )
            )

        # Upsert points to Qdrant
        self.client.upsert(
            collection_name=self.collection_name,
            points=points,
        )

        return [info.id for info in informations]

    def get_information(self, information_id: str) -> Optional[Information]:
        """
        Retrieve an information by its ID.

        Parameters
        ----------
        information_id : str
            The ID of the information to retrieve.

        Returns
        -------
        Optional[Information]
            The information, or None if not found.
        """
        points = self.client.retrieve(
            collection_name=self.collection_name,
            ids=[information_id],
        )

        if not points:
            return None

        point = points[0]
        return Information.from_base64(point.payload["info"])

    def similarity_search(
        self,
        query: Union[str, BaseModel, MediaBase, Information, List[float]],
        k: int = 5,
        threshold: float = 0.0,
    ) -> List[Information]:
        """
        Find similar informations using Qdrant's vector search.

        Parameters
        ----------
        query : Union[str, BaseModel, MediaBase, Information, List[float]]
            The query: a string, Pydantic model, media object, Information,
            or a raw embedding vector.
        k : int
            The maximum number of results to return.
        threshold : float
            Minimum similarity score. Results below this value are excluded.

        Returns
        -------
        List[Information]
            The most similar informations, ordered by descending similarity.
        """

        # If no informations are requested we obey:
        if k <= 0:
            return []

        if isinstance(query, list) and isinstance(query[0], float):

            # If the query is already a vector, we don't need to normalize it:
            query_vector = query

        else:
            # Normalize the query
            query = self._normalize_query(query)

            # Wrap in Information:
            query = Information(query)

            # Get the embedding for the query:
            query_vector = list(self.embedding_function([query])[0])

        # Convert to numpy array:
        query_vector = np.array(query_vector)

        # Normalize query vector
        query_norm = np.linalg.norm(query_vector)
        if query_norm > 0:
            query_vector = query_vector / query_norm

        # Perform similarity search
        search_result = self.client.query_points(
            collection_name=self.collection_name,
            query=query_vector,
            limit=k,
        )

        results = []
        for result in search_result.points:

            # Deserialize the information from the payload
            info = Information.from_base64(result.payload["info"])
            # Information.from_json(result.payload["info"])

            # Set the score (similarity) for the information
            info.score = result.score

            # Check if the score is above the threshold
            if info.score >= threshold:
                # Add the information to the list
                results.append(info)

        return results

    def delete(self, ids: List[str]) -> None:
        """
        Delete informations from the database.

        Parameters
        ----------
        ids: List[str]
            The IDs of the informations to delete.
        """
        self.client.delete(
            collection_name=self.collection_name,
            points_selector=ids,
        )

    def clear(self) -> None:
        """Clear all informations from the database."""
        try:
            # Delete the collection
            self.client.delete_collection(collection_name=self.collection_name)

            # Recreate the collection
            from qdrant_client.http import models

            self.client.create_collection(
                collection_name=self.collection_name,
                vectors_config=models.VectorParams(
                    size=self.embedding_length,
                    distance=models.Distance.COSINE,
                ),
            )
        except Exception as e:
            # If collection doesn't exist, that's fine
            aprint(f"Error clearing collection: {e}")

    # Implement AugmentationBase methods
    def get_relevant_informations(
        self,
        query: Union[str, BaseModel, MediaBase, Information],
        k: int = 5,
        threshold: float = 0.0,
    ) -> List[Information]:
        """
        Get relevant informations for a query via Qdrant similarity search.

        Parameters
        ----------
        query : Union[str, BaseModel, MediaBase, Information]
            The query to retrieve informations for.
        k : int
            The maximum number of informations to retrieve.
        threshold : float
            Minimum relevance score. Results below this value are excluded.

        Returns
        -------
        List[Information]
            A list of relevant informations sorted by similarity.
        """

        # Normalize the query
        query = self._normalize_query(query)

        # Perform similarity search
        return self.similarity_search(query, k, threshold)

    def save(self) -> None:
        """Persist the database state. No-op for Qdrant (auto-persisted)."""
        pass

    def close(self) -> None:
        """Close the connection to the Qdrant database."""
        if hasattr(self, "client") and self.client is not None:
            self.client.close()
