import pinecone
from typing import List
from .vectordb import VectorDB
from classes.secret import Secret


class Pinecone(VectorDB):
    def __init__(self, environment: str, index: str, api_key: Secret):
        pinecone.init(environment=environment, api_key=api_key)
        self.index = pinecone.Index(index)

    def delete_ns(self, ns: str):
        self.index.delete(delete_all=True, namespace=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]:
        return [
            {**x["metadata"], "_id": x["id"], "_score": x["score"]}
            for x in self.index.query(
                namespace=ns,
                top_k=limit,
                include_metadata=True,
                vector=vec,
            )["matches"]
        ]
    
    def upsert(self, ns: str, vec_id: str, vec: List[float], metadata: dict):
        self.index.upsert(
            vectors=[{
                "id": vec_id,
                "values": vec,
                "metadata": metadata,
            }],
            namespace=ns,
        )
        
    def delete(self, ns: str, vec_id: str):
        self.index.delete(ids=[vec_id], namespace=ns)