# app.py
import streamlit as st
import google.generativeai as genai
from langchain.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from io import StringIO

# Configure Gemini API
GEMINI_API_KEY = "YOUR_GEMINI_API_KEY"  # Replace with your Gemini API key
genai.configure(api_key=GEMINI_API_KEY)

# Function to load and index documents from an uploaded file
def load_and_index_documents(uploaded_file):
    # Read the uploaded file
    file_content = uploaded_file.read().decode("utf-8")
    
    # Use StringIO to simulate a file object for TextLoader
    file_like = StringIO(file_content)
    loader = TextLoader(file_like)
    documents = loader.load()
    
    # Split documents into chunks
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    texts = text_splitter.split_documents(documents)
    
    # Create embeddings and vector store
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    vector_store = FAISS.from_documents(texts, embeddings)
    return vector_store.as_retriever(search_kwargs={"k": 3})

# Generate answer using Gemini
def generate_answer(query, retriever):
    relevant_docs = retriever.get_relevant_documents(query)
    context = "\n".join([doc.page_content for doc in relevant_docs])
    prompt = f"""
    You are an AI tutor. Answer the following question based on the provided context:
    
    Question: {query}
    
    Context: {context}
    
    Answer:
    """
    model = genai.GenerativeModel('gemini-pro')
    response = model.generate_content(prompt)
    return response.text

# Streamlit app
st.title("AI Tutor")
st.write("Upload a file and ask me anything!")

# File uploader
uploaded_file = st.file_uploader("Upload a text or PDF file", type=["txt", "pdf"])

# Initialize retriever
retriever = None
if uploaded_file is not None:
    retriever = load_and_index_documents(uploaded_file)
    st.success("File uploaded and processed successfully!")

# Input from user
user_query = st.text_input("Enter your question:")

# Generate and display answer
if user_query and retriever:
    answer = generate_answer(user_query, retriever)
    st.write("Answer:")
    st.write(answer)
elif user_query and not retriever:
    st.warning("Please upload a file first!")