from PyPDF2 import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS
from transformers import pipeline
import streamlit as st
import requests
from io import BytesIO

# Set up Hugging Face model pipeline for text generation
pipe = pipeline("text-generation", model="meta-llama/Llama-Guard-3-8B-INT8")

# List of GitHub PDF URLs
PDF_URLS = [
    "https://github.com/TahirSher/GenAI_Lawyers_Guide/blob/main/bi%20pat%20graphs.pdf",
    "https://github.com/TahirSher/GenAI_Lawyers_Guide/blob/main/bi-partite.pdf",
    # Add more document links as needed
]

def fetch_pdf_text_from_github(urls):
    text = ""
    for url in urls:
        response = requests.get(url)
        if response.status_code == 200:
            pdf_file = BytesIO(response.content)
            pdf_reader = PdfReader(pdf_file)
            for page in pdf_reader.pages:
                page_text = page.extract_text()
                if page_text:
                    text += page_text
        else:
            st.error(f"Failed to fetch PDF from URL: {url}")
    return text

@st.cache_data
def get_text_chunks(text):
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=10000, chunk_overlap=1000)
    chunks = text_splitter.split_text(text)
    return chunks

@st.cache_resource
def load_or_create_vector_store(text_chunks):
    embeddings = FAISS.get_default_embeddings()
    vector_store = FAISS.from_texts(text_chunks, embedding=embeddings)
    return vector_store

def generate_answer(user_question, context_text):
    # Format the input message for the pipeline
    messages = [
        {"role": "user", "content": f"Context: {context_text}\nQuestion: {user_question}"}
    ]
    # Generate response using the pipeline
    response = pipe(messages, max_length=250, do_sample=True)
    return response[0]['generated_text'][:250]  # Limit response to 250 characters

def user_input(user_question, vector_store):
    docs = vector_store.similarity_search(user_question)
    context_text = " ".join([doc.page_content for doc in docs])
    return generate_answer(user_question, context_text)

def main():
    st.set_page_config(page_title="RAG-based PDF Chat", layout="centered", page_icon="📄")
    st.title("📄 Query PDF Documents on GitHub")

    # Load documents from GitHub
    raw_text = fetch_pdf_text_from_github(PDF_URLS)
    text_chunks = get_text_chunks(raw_text)
    vector_store = load_or_create_vector_store(text_chunks)

    # User question input
    user_question = st.text_input("Ask a Question:", placeholder="Type your question here...")

    if st.button("Get Response"):
        if not user_question:
            st.warning("Please enter a question before submitting.")
        else:
            with st.spinner("Generating response..."):
                answer = user_input(user_question, vector_store)
                st.markdown(f"**🤖 AI:** {answer}")

if __name__ == "__main__":
    main()