# -*- coding: utf-8 -*-

# ===============================================================
# osintgpt
#
# Author: @estebanpdl
#
# File: qdrant.py
# Description: Qdrant API. This file contains the Qdrant class
#   method for managing the Qdrant API connection.
# ===============================================================

# import modules <Qdrant>
import os
import sys
import qdrant_client

# import submodules <Qdrant>
from qdrant_client.http import models as rest

# import submodules
from dotenv import load_dotenv

# type hints
from typing import List, Optional

# import exceptions
from osintgpt.exceptions.errors import MissingEnvironmentVariableError

# import base class
from .base import BaseVectorEngine

# Qdrant class
class Qdrant(BaseVectorEngine):
    '''
    Qdrant class

    This class provides methods for managing connections to a Qdrant server,
    allowing users to store, retrieve, and manipulate high-dimensional vector
    embeddings and associated documents within a Qdrant collection. It offers
    functionality for creating, updating, and deleting collections, as well as
    adding and updating vector embeddings and their associated payloads.

    Main features:
        - Manage connections to a Qdrant server
        - Create, update, and delete Qdrant collections
        - Add and update vector embeddings and payloads
        - Retrieve collection information and count vectors
        - Efficiently store and search embeddings

    For more information about QdrantClient arguments, see:
    github.com/qdrant/qdrant-client/blob/master/qdrant_client/qdrant_client.py
    '''
    # constructor
    def __init__(self, env_file_path: str):
        '''
        Constructor

        args:
            **kwargs: keyword arguments for QdrantClient
        '''
        # load environment variables
        load_dotenv(dotenv_path=env_file_path)

        # set environment variables
        self.set_required_variables()

    # set required environment variables
    def set_required_variables(self):
        '''
        set required environment variables

        This method sets for the required environment variables for connecting
        to a Qdrant server.

        returns:
            use_remote: use remote
            use_local: use local
        '''
        # set required environment variables
        use_remote = os.getenv('QDRANT_API_KEY') and os.getenv('QDRANT_URL')
        use_local = os.getenv('QDRANT_PORT') and os.getenv('QDRANT_HOST')

        if not (use_remote or use_local):
            raise MissingEnvironmentVariableError(
                'QDRANT_API_KEY or QDRANT_URL or QDRANT_HOST or QDRANT_PORT'
            )

        # set environment variables
        if use_remote:
            self.api_key = os.getenv('QDRANT_API_KEY')
            self.url = os.getenv('QDRANT_URL')

            # connect
            self.qdrant = qdrant_client.QdrantClient(
                url=self.url,
                api_key=self.api_key,
                https=True
            )
        else:
            self.host = os.getenv('QDRANT_HOST')
            self.port = int(os.getenv('QDRANT_PORT'))

            # connect
            self.qdrant = qdrant_client.QdrantClient(
                host=self.host,
                port=self.port
            )
        
        '''

        Ensure if is indeed connected
        '''
        # Perform a simple operation to check connectivity
        try:
            collections = self.get_collections()
        except Exception as e:
            m = f'''
            Unable to establish a connection to the Qdrant server. Please ensure
            that the Qdrant server is up and running. If you're using this locally,
            make sure to start the Qdrant server before using this feature.
            '''
            raise ConnectionError(' '.join(m.split()).strip()) from None
    
    # get client
    def get_client(self):
        '''
        Get client

        returns:
            client: qdrant client
        '''
        # return client
        return self.qdrant
    
    # get collections
    def get_collections(self):
        '''
        Get collections

        returns:
            collections: collections
        '''
        # get collections
        collections = self.qdrant.get_collections()
        
        # return collections
        return collections
    
    # get collection info
    def get_collection_info(self, collection_name: str):
        '''
        Get collection info

        args:
            collection_name: collection name
                type: str
        
        returns:
            info: collection info
        '''
        # get collection
        info = self.qdrant.get_collection(collection_name=collection_name)
        
        # return collection info
        return info
    
    # count vectors
    def count_vectors(self, collection_name: str):
        '''
        Count vectors

        args:
            collection_name: collection name
                type: str
        
        returns:
            count: count
        '''
        # return count
        return self.qdrant.count(collection_name=collection_name)
    
    def _validate_payload_length(self, payload: Optional[List], vectors: List):
        '''
        Validate payload length

        args:
            payload: payload
                type: list
            vectors: vectors
                type: list
        '''
        # validate payload length
        if payload and len(payload) != len(vectors):
            raise ValueError('Payload length must be the same as vectors length')
    
    # get vector
    def get_vectors(self, collection_name: str):
        '''
        Get vector

        args:
            collection_name: collection name
                type: str
        
        returns:
            vector: vector
        '''
        # get collection
        collection = self.qdrant.get_collection(collection_name=collection_name)
        
        # get vector
        vector = collection.config.params
        return vector
    
    # create collection
    def create_collection(self, collection_name: str, vector_size: int,
        vector_name: str = 'main'):
        '''
        Create collection

        args:
            collection_name: collection name
                type: str
            vector_size: vector size
                type: int
            vector_name: name
                type: str
        '''
        # create collection
        self.qdrant.recreate_collection(
            collection_name=collection_name,
            vectors_config={
                vector_name: rest.VectorParams(
                    distance=rest.Distance.COSINE,
                    size=vector_size
                )
            }
        )
    
    # add vectors
    def add_vectors(self, collection_name: str, vectors: List,
        vector_name: str = 'main', payload: Optional[List[dict]] = None):
        '''
        Add vectors

        args:
            collection_name: collection name
                type: str
            vectors: vectors
                type: list
            vector_name: name
                type: str
            payload: payload. Should be the same length as vectors and same order.
                type: list
        '''
        # validate payload length
        self._validate_payload_length(payload, vectors)
        
        # add vectors
        self.qdrant.upsert(
            collection_name=collection_name,
            points=[
                rest.PointStruct(
                    id=k,
                    vector={
                        vector_name: v
                    },
                    payload=payload[k] if payload else None
                )
                for k, v in enumerate(vectors)
            ]
        )
    
    # update vector collection
    def update_vector_collection(self, collection_name: str, vectors: List,
        vector_name: str = 'main', payload: Optional[List[dict]] = None):
        '''
        Update vector collection

        args:
            collection_name: collection name
                type: str
            vectors: vectors
                type: list
            vector_name: name
                type: str
            payload: payload. Should be the same length as vectors and same order.
                type: List[dict]
        '''
        # validate payload length
        self._validate_payload_length(payload, vectors)

        # count vectors
        count = self.count_vectors(collection_name=collection_name)
        n = count.count

        # add vectors
        self.qdrant.upsert(
            collection_name=collection_name,
            points=[
                rest.PointStruct(
                    id=k + n,
                    vector={
                        vector_name: v
                    },
                    payload=payload[k] if payload else None
                )
                for k, v in enumerate(vectors)
            ]
        )
    
    # delete collection
    def delete_collection(self, collection_name: str):
        '''
        Delete collection

        args:
            collection_name: collection name
                type: str
        '''
        # delete collection
        self.qdrant.delete_collection(collection_name=collection_name)
    
    # search query
    def search_query(self, embedded_query: List[float], top_k: int = 10, **kwargs):
        '''
        Search query in collection

        args:
            embedded_query: embedded query
                type: list
            top_k: top k
                type: int

            kwargs:
                collection_name: collection name
                    type: str
                vector_name: name
                    type: str
        
        returns:
            result: result
        '''
        # collection name
        collection_name = kwargs.get('collection_name', None)
        if collection_name is None:
            raise ValueError('collection_name must be specified')
        
        # vector name
        vector_name = kwargs.get('vector_name', 'main')

        # query results
        query_results = self.qdrant.search(
            collection_name=collection_name,
            query_vector=(
                vector_name, embedded_query
            ),
            limit=top_k
        )

        return query_results
