from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

from src.embeddings.embedding_factory import get_text_embedding
from src.retrieval.vector_store import VectorStoreFactory
from src.llm.llm_factory import get_llm
from src.utils.logger import get_logger

logger = get_logger(__name__)


def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)


def main():

    logger.info("Starting RAG query interface...")

    embedding = get_text_embedding()
    vectordb = VectorStoreFactory.create(embedding)
    retriever = vectordb.as_retriever(search_kwargs={"k": 5})

    llm = get_llm()

    prompt = ChatPromptTemplate.from_template("""
You are an anatomy tutor.

Answer the question using ONLY the context below.

Context:
{context}

Question:
{question}
""")

    rag_chain = (
        prompt
        | llm
        | StrOutputParser()
    )


    


    while True:

        query = input("\nAsk a question (or type 'exit'): ")

        if query.lower() == "exit":
            break

        # 🔎 STEP 1 — RETRIEVE DOCUMENTS
        docs = retriever.invoke(query)

        # 🧪 DEBUG: SEE WHAT IS RETRIEVED (keep for now)
        print("\nRETRIEVED CHUNKS:\n")
        for d in docs:
            print(d.page_content[:300])
            print("------")

        # 🧠 STEP 2 — PREPARE CONTEXT FOR LLM
        context = "\n\n".join(doc.page_content for doc in docs)

        response = rag_chain.invoke({
            "context": context,
            "question": query
        })

        # 🧾 STEP 3 — PRINT ANSWER
        print("\nANSWER:\n")
        print(response)

        # 📚 STEP 4 — PRINT SOURCES
        print("\nSOURCES:\n")
        for doc in docs:
            source = doc.metadata.get("source", "unknown")
            page = doc.metadata.get("page", "unknown")
            print(f"{source} — page {page}")


if __name__ == "__main__":
    main()