# src/app.py
import gradio as gr
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain_community.llms import HuggingFacePipeline
from retriever import get_retriever
from transformers import pipeline
import os

# Load local HuggingFace pipeline
pipe = pipeline(
    "text-generation",
    model="tiiuae/falcon-7b-instruct",
    trust_remote_code=True,
    device_map="auto",
    max_new_tokens=512,
    temperature=0.2
)
llm = HuggingFacePipeline(pipeline=pipe)

retriever = get_retriever()

# Prompt template
template = """
You are a legal assistant. Use the provided context to answer the question.
If language mode is Nigerian Pidgin, respond in Nigerian Pidgin.

Question: {question}

Context:
{context}

Answer:
"""

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

qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=retriever,
    chain_type="stuff",
    return_source_documents=True,  # Needed to list references
    chain_type_kwargs={"prompt": prompt}
)

def answer_question(user_input, lang_choice):
    if lang_choice == "pidgin":
        user_input = f"Respond in Nigerian Pidgin: {user_input}"

    result = qa_chain(user_input)
    answer_text = result["result"]

    # Collect unique source file names
    sources = list({doc.metadata.get("source", "Unknown") for doc in result["source_documents"]})
    sources_list = "\n".join(f"- {src}" for src in sources)

    return f"{answer_text}\n\nReferences:\n{sources_list}"

def launch_interface():
    iface = gr.Interface(
        fn=answer_question,
        inputs=[
            gr.Textbox(label="Your question"),
            gr.Radio(["english", "pidgin"], label="Language")
        ],
        outputs=gr.Textbox(label="Answer"),
        title="KnowYourRight Bot",
        description="Ask legal rights questions in English or Nigerian Pidgin with references"
    )
    iface.launch()

if __name__ == "__main__":
    launch_interface()