import os
import requests
import streamlit as st
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.document_loaders import PyPDFLoader
from langchain.prompts import PromptTemplate
from langchain.chains import RetrievalQA
from groq import Groq

# Initialize Groq client
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))


class GroqLLM:
    def _call(self, prompt: str, stop=None):
        response = client.chat.completions.create(
            messages=[
                {"role": "user", "content": prompt},
            ],
            model="llama-3.3-70b-versatile",
        )
        return response.choices[0].message.content

    @property
    def _llm_type(self) -> str:
        return "Groq"


# Title of the application
st.title("PPRA Rules 2004 - Q&A Application")

# Download and load the PDF
pdf_url = "https://drive.google.com/uc?id=1faNpSV_UIZzd3h08qtzvSRGmzDkNtmuA"
pdf_response = requests.get(pdf_url)
pdf_path = "ppra_rules_2004.pdf"

with open(pdf_path, "wb") as f:
    f.write(pdf_response.content)

# Load the PDF document
loader = PyPDFLoader(pdf_path)
documents = loader.load()

# Initialize the embeddings and vectorstore
embeddings = HuggingFaceEmbeddings()
vectorstore = FAISS.from_documents(documents, embeddings)
retriever = vectorstore.as_retriever()

# Custom prompt template
prompt_template = """
You are an AI assistant tasked with answering questions about the Public Procurement Rules, 2004 (PPRA Rules) in Pakistan.
Use the provided context to answer the user's question as accurately as possible.

Context:
{context}

Question:
{query}

Answer:
"""
qa_prompt = PromptTemplate(input_variables=["context", "query"], template=prompt_template)

# Create the QA chain
qa_chain = RetrievalQA.from_chain_type(
    llm=GroqLLM(),
    retriever=retriever,
    combine_documents_chain_kwargs={"prompt": qa_prompt},
    return_source_documents=True,
)

# User interaction
st.subheader("Ask Questions About PPRA Rules 2004")
user_query = st.text_input("Enter your question:")

if user_query:
    try:
        # Run the query through the QA chain
        response = qa_chain({"query": user_query})

        # Display the answer
        st.subheader("Answer:")
        st.write(response["result"])

        # Display the relevant sources
        st.subheader("Relevant Sources:")
        for doc in response["source_documents"]:
            st.write(doc.page_content[:500])  # Show the first 500 characters
    except Exception as e:
        st.error(f"An error occurred: {str(e)}")