Update app.py

app.py

@@ -1,59 +1,98 @@
 import streamlit as st
-from config import process_pdf, process_image, process_audio, generate_flashcards, export_to_apkg
+from config import (
+    process_pdf, process_image, process_audio, process_text, process_youtube,
+    generate_flashcards, export_to_apkg, export_to_csv
+)
 import tempfile
 import os
+import uuid
 
-st.set_page_config(page_title="🦇 BatAnki", layout="wide")
+st.set_page_config(page_title="🦇 BatAnki Pro", layout="wide")
+st.title("🦇 BatAnki – AI Flashcard Engine for Learners")
 
-st.title("🦇 BatAnki – AI Flashcard Generator")
-st.markdown("Built for med learners. Powered by AI. Inspired by Batman.")
+st.sidebar.title("📁 Upload Source")
+upload_type = st.sidebar.radio("Select input type", ["PDF", "Image", "Audio", "Text File", "YouTube Link"])
+deck_name = st.sidebar.text_input("📌 Deck name", value="BatAnki_Deck")
+model_choice = st.sidebar.selectbox("🧠 AI Model", ["t5-base", "microsoft/BioGPT"])
+generate_types = st.sidebar.multiselect("Card Types", ["Q&A", "Cloze", "MCQ", "Reverse"], default=["Q&A"])
 
-# Sidebar
-st.sidebar.title("📁 Upload Your Material")
-upload_type = st.sidebar.radio("Select input type:", ["PDF", "Image", "Audio"])
+uploaded_file = None
+yt_url = ""
 
-uploaded_file = st.sidebar.file_uploader("Upload your file", type=["pdf", "png", "jpg", "jpeg", "mp3", "wav"])
+if upload_type == "YouTube Link":
+    yt_url = st.sidebar.text_input("📺 YouTube URL")
+else:
+    uploaded_file = st.sidebar.file_uploader("Upload file", type=["pdf", "png", "jpg", "jpeg", "mp3", "wav", "txt", "docx"])
 
-if uploaded_file:
-    with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
-        tmp_file.write(uploaded_file.read())
-        file_path = tmp_file.name
+# Extract text from the file
+extracted_text = ""
+file_path = None
 
-    st.success("✅ File uploaded successfully!")
+if uploaded_file:
+    with tempfile.NamedTemporaryFile(delete=False) as tmp:
+        tmp.write(uploaded_file.read())
+        file_path = tmp.name
 
+if st.sidebar.button("🧠 Extract Text"):
     if upload_type == "PDF":
         extracted_text = process_pdf(file_path)
     elif upload_type == "Image":
         extracted_text = process_image(file_path)
     elif upload_type == "Audio":
         extracted_text = process_audio(file_path)
+    elif upload_type == "Text File":
+        extracted_text = process_text(file_path)
+    elif upload_type == "YouTube Link" and yt_url:
+        extracted_text = process_youtube(yt_url)
     else:
-        extracted_text = ""
-
-    if extracted_text:
-        st.subheader("📝 Extracted Text Preview")
-        st.text_area("Text", extracted_text[:5000], height=300)
-
-        st.subheader("🧠 Generate Flashcards")
-        if st.button("Generate"):
-            with st.spinner("Generating cards using AI..."):
-                cards = generate_flashcards(extracted_text)
-            st.success(f"✅ Generated {len(cards)} cards!")
-
-            for i, card in enumerate(cards, 1):
-                st.markdown(f"**Card {i}**")
-                st.markdown(f"- **Q:** {card['question']}")
-                st.markdown(f"- **A:** {card['answer']}")
-                st.markdown("---")
-
-            st.subheader("📦 Export Deck")
-            deck_name = st.text_input("Enter deck name", value="BatAnki_Deck")
-            if st.button("Export to APKG"):
-                apkg_path = export_to_apkg(cards, deck_name)
-                with open(apkg_path, "rb") as f:
-                    st.download_button("⬇️ Download .apkg", f, file_name=f"{deck_name}.apkg")
-
-        else:
-            st.info("Click 'Generate' to create flashcards.")
-
-    os.remove(file_path)
+        st.warning("Please upload or enter a valid input.")
+
+    if not extracted_text.strip():
+        st.error("No text extracted.")
+    else:
+        st.success("✅ Text extracted!")
+        st.session_state["extracted_text"] = extracted_text
+        st.session_state["deck_name"] = deck_name
+        st.session_state["generate_types"] = generate_types
+        st.session_state["model"] = model_choice
+
+if "extracted_text" in st.session_state:
+    st.subheader("📄 Extracted Text")
+    st.text_area("Extracted Content", st.session_state["extracted_text"][:10000], height=300)
+
+    if st.button("🚀 Generate Flashcards"):
+        with st.spinner("Generating cards..."):
+            cards = generate_flashcards(
+                st.session_state["extracted_text"],
+                model_name=st.session_state["model"],
+                types=st.session_state["generate_types"]
+            )
+            st.session_state["cards"] = cards
+        st.success(f"✅ Generated {len(cards)} cards")
+
+if "cards" in st.session_state:
+    st.subheader("🧾 Flashcard Preview & Tagging")
+    tagged_cards = []
+    for idx, card in enumerate(st.session_state["cards"], 1):
+        st.markdown(f"**Card {idx}**")
+        question = st.text_input(f"Q{idx}", card["question"])
+        answer = st.text_area(f"A{idx}", card["answer"])
+        tag = st.text_input(f"🏷️ Tag for Card {idx}", value=card.get("tag", ""))
+        tagged_cards.append({"question": question, "answer": answer, "tag": tag})
+        st.markdown("---")
+    st.session_state["cards"] = tagged_cards
+
+    st.subheader("📦 Export")
+    col1, col2 = st.columns(2)
+
+    with col1:
+        if st.button("📥 Export to Anki (.apkg)"):
+            path = export_to_apkg(st.session_state["cards"], st.session_state["deck_name"])
+            with open(path, "rb") as f:
+                st.download_button("⬇️ Download APKG", f, f"{st.session_state['deck_name']}.apkg")
+
+    with col2:
+        if st.button("📊 Export to CSV"):
+            path = export_to_csv(st.session_state["cards"], st.session_state["deck_name"])
+            with open(path, "rb") as f:
+                st.download_button("⬇️ Download CSV", f, f"{st.session_state['deck_name']}.csv")