from typing import List
from .vectordb import VectorDB
from classes.secret import Secret
from qdrant_client import QdrantClient
from qdrant_client.models import PointStruct, Distance, VectorParams, PointIdsList


class Qdrant(VectorDB):
    def __init__(self, url: str, api_key: Secret = None):
        self.client = QdrantClient(url=url, api_key=api_key)

    def create_ns(self, ns: str, size: int):
        self.client.recreate_collection(
            collection_name=ns,
            vectors_config=VectorParams(size=size, distance=Distance.COSINE),
        )
        
    def delete_ns(self, ns: str):
        self.client.delete_collection(collection_name=ns)
        
    def vec_id(self, index: List[str]) -> str:
        h = hashlib.new('md5')
        for x in index:
            h.update(hashlib.sha256(x.encode()).digest())
        return h.hexdigest()

    def retrieve(self, ns: str, vec: List[float], limit: int) -> List[dict]:
        xs = self.client.search(
            collection_name=ns,
            query_vector=vec,
            limit=limit,
        )
        return [{
            **x.payload,
            '_id': x.id,
            '_score': x.score,
        } for x in xs]

    def upsert(self, ns: str, vec_id: str, vec: List[float], metadata: dict):
        self.client.upsert(
            collection_name=ns,
            points=[
                PointStruct(
                    id=vec_id,
                    vector=vec,
                    payload=metadata,
                )
            ]
        )

    def delete(self, ns: str, vec_id: str):
        self.client.delete(
            collection_name=ns,
            points_selector=PointIdsList(points=[vec_id]),
        )