from typing import List, Dict, Union
import datetime
import asyncio
from pinecone import Pinecone
from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.vector_stores.pinecone import PineconeVectorStore

class PineconeIndex:
    def __init__(self, api_key: str, index_name: str, index_namespace: str):
        self._index_name = index_name
        self._index_namespace = index_namespace
        self._pc = Pinecone(api_key=api_key)
        self.pc_index = self._set_index(index_name, index_namespace)
    
    def _set_index(self, index_name: str, index_namespace: str) -> VectorStoreIndex:
        vector_store = PineconeVectorStore(
            pinecone_index=self._pc.Index(index_name),
            add_sparse_vector=True, 
            namespace=index_namespace
        )
        storage_context = StorageContext.from_defaults(vector_store=vector_store)
        pc_index = VectorStoreIndex.from_vector_store(vector_store=vector_store, storage_context=storage_context)
        return pc_index

    def retrieve_context(self, query: str) -> Dict[str, Union[str, int]]:
        start_time = round(datetime.datetime.now().timestamp() * 1000)
        response = self.pc_index.as_query_engine(similarity_top_k=3).query(query)
        end_time = round(datetime.datetime.now().timestamp() * 1000)
        return {"result": response.response, "start": start_time, "end": end_time}

    async def aretrieve_context(self, query: str) -> Dict[str, Union[str, int]]:
        start_time = round(datetime.datetime.now().timestamp() * 1000)
        response = await self.pc_index.as_query_engine(
            similarity_top_k=3,
            use_async=True
        ).aquery(query)
        end_time = round(datetime.datetime.now().timestamp() * 1000)
        return {"result": response.response, "start": start_time, "end": end_time}

    async def aretrieve_context_multi(self, query_list: List[str]) -> List[Dict]:
        result = await asyncio.gather(*(self.aretrieve_context(query) for query in query_list))
        return result