from .config import Config
from pinecone import Pinecone
from typing import List, Optional


class VectorStore:
    def __init__(self):
        if Config.VECTOR_STORE_TYPE == "pinecone":
            self.pc = Pinecone(api_key=Config.PINECONE_API_KEY)
            self.index = self.pc.Index(Config.PINECONE_INDEX)

    def upsert_document(self, doc_id: str, embedding: List[float], metadata: dict):
        self.index.upsert(
            vectors=[{
                "id": doc_id,
                "values": embedding,
                "metadata": metadata
            }]
        )

    def search_similar(self, embedding: List[float], top_k: int = 3):
        return self.index.query(
            vector=embedding,
            top_k=top_k,
            include_metadata=True
        )

    def get_document(self, doc_id: str) -> Optional[dict]:
        fetch_response = self.index.fetch(ids=[doc_id])
        if doc_id in fetch_response.vectors:
            return fetch_response.vectors[doc_id]
        return None