import streamlit as st
from model.questionAnsweringBot import QuestionAnsweringBot
from model.retriever import Retriever

def process_query(llm_key, query, retrieval_method):
    if "retriever" not in st.session_state:
        st.session_state.retriever = Retriever()
        print("Loading and preparing dataset...")
        st.session_state.retriever.load_and_prepare_dataset()
        st.session_state.retriever.prepare_bm25()
        st.session_state.retriever.compute_embeddings()

    retriever = st.session_state.retriever

    if retrieval_method == "BM25":
        print("Retrieving documents using BM25...")
        retrieved_docs = retriever.retrieve_documents_bm25(query)
    else:
        print("Retrieving documents using Semantic Search...")
        retrieved_docs = retriever.retrieve_documents_semantic(query)

    bot = QuestionAnsweringBot(llm_key)
    prompt = getPrompt(retrieved_docs, query)
    answer = bot.generate_answer(prompt)

    return retrieved_docs, answer

def getPrompt(retrieved_docs, query):
    prompt = (
        "You are an LM integrated into an RAG system that answers questions based on provided documents.\n"
        "Rules:\n"
        "- Reply with the answer only and nothing but the answer.\n"
        "- Say 'I don't know' if you don't know the answer.\n"
        "- Use only the provided documents.\n"
        "- Citations are required. Include the document and chunk number in square brackets after the information (e.g., [Document 1, Chunk 2]).\n\n"
        "Documents:\n"
    )

    for i, doc in enumerate(retrieved_docs):
        prompt += f"Document {i + 1}: {doc}\n"

    prompt += f"\nQuery: {query}\n"

    return prompt