import os
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_community.vectorstores import FAISS

LLM_PROVIDER = os.getenv("LLM_PROVIDER", "groq")
GROQ_MODEL = os.getenv("GROQ_MODEL", "llama-3.1-8b-instant")

RAG_PROMPT_TEMPLATE = """You are a helpful assistant that answers questions based only on the provided context.

If the answer is not found in the context, say:
"I don't have enough information in the uploaded documents to answer that question."

Context:
{context}

Question: {question}

Answer:"""

PROMPT = PromptTemplate(
    template=RAG_PROMPT_TEMPLATE,
    input_variables=["context", "question"],
)

def _get_llm():
    if LLM_PROVIDER == "groq":
        api_key = os.getenv("GROQ_API_KEY")
        if not api_key:
            raise EnvironmentError("GROQ_API_KEY is not set in Hugging Face Secrets.")
        from langchain_groq import ChatGroq
        print(f"  Using Groq model: {GROQ_MODEL}")
        return ChatGroq(model=GROQ_MODEL, api_key=api_key, temperature=0)
    else:
        raise ValueError(f"Unknown LLM_PROVIDER: '{LLM_PROVIDER}'.")

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

def build_rag_chain(vector_store, k=4):
    llm = _get_llm()
    retriever = vector_store.as_retriever(
        search_type="similarity",
        search_kwargs={"k": k},
    )
    chain = (
        {"context": retriever | _format_docs, "question": RunnablePassthrough()}
        | PROMPT
        | llm
        | StrOutputParser()
    )
    print("  RAG chain ready.")
    return chain, retriever

def ask_question(chain_tuple, question: str) -> dict:
    if not question.strip():
        return {"answer": "Please enter a question.", "sources": []}
    chain, retriever = chain_tuple
    answer = chain.invoke(question)
    sources = retriever.invoke(question)
    return {
        "answer": answer,
        "sources": sources,
    }