import streamlit as st
from utils.pdf_parser import extract_text_from_pdf
from utils.ocr import extract_text_from_image
from utils.memory_manager import save_to_memory, retrieve_memory
from utils.quiz_generator import generate_quiz
from utils.flashcards import generate_flashcards
import os

st.set_page_config(page_title="AI Learning Companion", layout="wide")
st.title("📚 AI Learning Companion with Memory")

uploaded_file = st.file_uploader("Upload Notes (PDF or Image)", type=["pdf", "png", "jpg"])

if uploaded_file:
    file_path = os.path.join("data", uploaded_file.name)
    with open(file_path, "wb") as f:
        f.write(uploaded_file.getbuffer())

    if uploaded_file.name.endswith(".pdf"):
        extracted_text = extract_text_from_pdf(file_path)
    else:
        extracted_text = extract_text_from_image(file_path)

    st.subheader("Extracted Text")
    st.write(extracted_text)

    if st.button("Save to Memory"):
        save_to_memory(extracted_text, {"filename": uploaded_file.name})
        st.success("Saved to memory!")

st.subheader("Search Your Memory")
search_query = st.text_input("Enter a topic to search")
if search_query:
    results = retrieve_memory(search_query)
    st.write(results)

st.subheader("Generate Study Material")
topic_text = st.text_area("Enter text or select from memory results")
if st.button("Generate Quiz"):
    st.write(generate_quiz(topic_text))

if st.button("Generate Flashcards"):
    st.write(generate_flashcards(topic_text))