from dotenv import load_dotenv
import os
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.huggingface import HuggingFaceInferenceAPI
from llama_index.core.settings import Settings
from llama_index.core import SimpleDirectoryReader
from llama_index.core.node_parser import SentenceWindowNodeParser
from pinecone import Pinecone
from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.vector_stores.pinecone import PineconeVectorStore
from llama_index.core.postprocessor import MetadataReplacementPostProcessor
from llama_index.core.postprocessor import SentenceTransformerRerank


def initialize_llm():
    load_dotenv()
    HF_TOKEN = os.getenv("HUGGINGFACE_API_KEY")
    
    Settings.llm = HuggingFaceInferenceAPI(model_name="mistralai/Mixtral-8x7B-Instruct-v0.1", model_kwargs={"temperature": 0.1}, token=HF_TOKEN)
    Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")




def initialize_pinecone():
    load_dotenv()
    api_key = os.environ["PINECONE_API_KEY"]
    index_name = "chatbot"
    pc = Pinecone(api_key=api_key)
    pinecone_index = pc.Index(index_name)
    vector_store = PineconeVectorStore(pinecone_index=pinecone_index)
    index = VectorStoreIndex.from_vector_store(vector_store=vector_store)
    return index

def create_query_engine(index):
    postproc = MetadataReplacementPostProcessor(target_metadata_key="window")
    rerank = SentenceTransformerRerank(top_n=2, model="BAAI/bge-reranker-base")
    query_engine = index.as_query_engine(similarity_top_k = 6,vector_store_query_mode="hybrid",alpha=0.5,node_postprocessors = [postproc, rerank],
)
    return query_engine

def generate_prompt(context, query):
    prompt_template = """
    You are a highly informed representative of KYC SUD CONSULTING. Your role is to assist users by answering questions about our company with accurate and helpful information. Use the provided context to guide your responses effectively. If a question falls outside your knowledge or the context provided, simply state that you don't have the information.

    Given the following context:
    {context_str}

    Answer the following question in different words while keeping the meaning intact:
    {query_str}

    Response:
    """
    return prompt_template.format(context_str=context, query_str=query)



# Function to process input sentence and return response
def get_response(query_engine, sentence):
    retrieved_context = query_engine.query(sentence)
    prompt = generate_prompt(retrieved_context, sentence)
    response = Settings.llm.complete(prompt)  # Use the correct method to generate response
    return response.text