import gradio as gr
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from transformers import pipeline
from PyPDF2 import PdfReader

# Load AI models
embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")  # Converts text to numbers
llm_pipeline = pipeline("text2text-generation", model="google/flan-t5-small")  # AI that answers questions

# Memory (Database) to store text
index = None
chunks = []

# Load and process document
def load_document(file):
    global index, chunks
    text = ""
    
    # Determine file path or object
    # If file is a string, it's a file path
    if isinstance(file, str):
        file_path = file
    else:
        # If file is not a string, try to use its .name attribute
        file_path = file.name

    # Read PDF or text file
    if file_path.endswith(".pdf"):
        reader = PdfReader(file_path)
        text = "\n".join(
            [page.extract_text() for page in reader.pages if page.extract_text()]
        )
    else:
        with open(file_path, "r", encoding="utf-8") as f:
            text = f.read()

    # Break text into small parts
    sentences = text.split(". ")
    chunks = [" ".join(sentences[i:i + 5]) for i in range(0, len(sentences), 5)]

    # Create embeddings and store in FAISS
    embeddings = np.array([embedding_model.encode(chunk) for chunk in chunks])
    index = faiss.IndexFlatL2(embeddings.shape[1])
    index.add(embeddings)

    return "📄 Document is ready! Now ask your question."

# Find and answer questions
def get_answer(query):
    if index is None:
        return "❌ Please upload a document first."
    
    # Find best matching text
    query_embedding = embedding_model.encode(query).reshape(1, -1)
    distances, indices = index.search(query_embedding, 3)
    retrieved_text = " ".join([chunks[i] for i in indices[0]])

    # Ask AI to generate an answer
    input_text = f"Question: {query}\nContext: {retrieved_text}"
    response = llm_pipeline(input_text, max_length=100)[0]['generated_text']
    
    return response

# Webpage design
with gr.Blocks() as demo:
    gr.Markdown("# 📚 Smart Study Helper")
    
    file_input = gr.File(label="Upload a textbook or notes (PDF/TXT)", file_types=[".pdf", ".txt"])
    upload_button = gr.Button("Process Document")
    status_text = gr.Textbox(label="📢 Status", interactive=False)

    query_input = gr.Textbox(label="Ask a question from the document:")
    query_button = gr.Button("Get Answer")
    output_text = gr.Textbox(label="🤖 AI Answer", interactive=False)

    upload_button.click(load_document, inputs=file_input, outputs=status_text)
    query_button.click(get_answer, inputs=query_input, outputs=output_text)

# Run the app
demo.launch()