from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.documents import Document
from dotenv import load_dotenv
import os
from src.config import GROQ_API_KEY, LLM_MODEL, LLM_TEMPERATURE, RETRIEVAL_K

load_dotenv()

def ensure_groq_key():
    if not GROQ_API_KEY:
        raise RuntimeError("GROQ_API_KEY is not configured. Set it in .env or Hugging Face secrets.")


def format_docs(docs: list[Document]) -> str:    return "\n\n".join(doc.page_content for doc in docs)

def build_qa_chain(vectorstore):
    ensure_groq_key()
    llm = ChatGroq(
        model=LLM_MODEL,
        api_key=GROQ_API_KEY,
        temperature=LLM_TEMPERATURE
    )

    prompt = ChatPromptTemplate.from_template("""You are a research assistant. Use the context below to answer the question accurately and concisely.
If the answer is not in the context, say "I don't have enough information in the provided document."

Context: {context}
Question: {question}

Answer:""")

    retriever = vectorstore.as_retriever(search_kwargs={"k": RETRIEVAL_K})

    chain = (
        {
            "context": lambda x: format_docs(retriever.invoke(x["question"])),
            "question": lambda x: x["question"]
        }
        | prompt
        | llm
        | StrOutputParser()
    )

    return chain

def get_answer_with_sources(vectorstore, question: str) -> tuple[str, list]:
    ensure_groq_key()
    retriever = vectorstore.as_retriever(search_kwargs={"k": RETRIEVAL_K})
    docs = retriever.invoke(question)
    context = format_docs(docs)

    llm = ChatGroq(
        model=LLM_MODEL,
        api_key=GROQ_API_KEY,
        temperature=LLM_TEMPERATURE
    )

    prompt = ChatPromptTemplate.from_template("""You are a research assistant. Use the context below to answer the question accurately and concisely.
If the answer is not in the context, say "I don't have enough information in the provided document."

Context: {context}
Question: {question}

Answer:""")

    chain = prompt | llm | StrOutputParser()
    answer = chain.invoke({"context": context, "question": question})
    sources = [doc.page_content[:200] for doc in docs]

    return answer, sources