import streamlit as st
import os
import tempfile
import fitz
from config import (
    process_pdf,
    process_text,
    process_image,
    process_audio,
    process_youtube,
    generate_flashcards,
    export_to_apkg,
    export_to_csv,
)

st.set_page_config(page_title="🦇 BatAnki AI", layout="wide")
st.title("📚 BatAnki – AI Flashcard Generator")

st.sidebar.markdown("📁 **Input Options**")
uploaded_file = st.sidebar.file_uploader("Upload a file", type=["pdf", "txt", "docx", "jpg", "png", "mp3", "wav"])
youtube_url = st.sidebar.text_input("Or paste YouTube URL")

deck_name = st.text_input("Deck Name", value="BatAnkiDeck")
types_selected = st.multiselect("Card Types", ["Q&A", "Cloze", "MCQ", "Reverse"], default=["Q&A"])
max_cards = st.slider("Max Cards", 5, 1000, 50)

input_text = ""
cards = []

if uploaded_file:
    suffix = uploaded_file.name.split(".")[-1]
    with tempfile.NamedTemporaryFile(delete=False, suffix=f".{suffix}") as tmp_file:
        tmp_file.write(uploaded_file.read())
        tmp_path = tmp_file.name

    if suffix == "pdf":
        doc = fitz.open(tmp_path)
        st.info("PDF preview:")
        page_number = st.number_input("Page", 1, len(doc), 1)
        page = doc[page_number - 1]
        pix = page.get_pixmap()
        st.image(pix.tobytes("png"), caption=f"Page {page_number}")
        text = page.get_text()
        input_text = text if text.strip() else process_pdf(tmp_path)
        if st.button("Generate Cards from This Page"):
            cards = generate_flashcards(input_text, types_selected, max_cards)
    elif suffix in ["jpg", "png"]:
        input_text = process_image(tmp_path)
        st.image(tmp_path)
    elif suffix in ["mp3", "wav"]:
        input_text = process_audio(tmp_path)
    elif suffix in ["txt", "docx"]:
        input_text = process_text(tmp_path)

elif youtube_url:
    st.info("Processing YouTube audio...")
    input_text = process_youtube(youtube_url)

if input_text and not cards:
    if st.button("Generate Cards"):
        cards = generate_flashcards(input_text, types_selected, max_cards)

if cards:
    st.subheader("Generated Flashcards")
    for i, card in enumerate(cards):
        st.markdown(f"**{i+1}. {card['question']}**")
        st.markdown(f"*Answer:* {card['answer']}")
        st.markdown("---")

    col1, col2 = st.columns(2)
    with col1:
        if st.button("Export to CSV"):
            export_to_csv(cards)
            st.success("Exported CSV")
    with col2:
        if st.button("Export to Anki (.apkg)"):
            path = export_to_apkg(cards, deck_name)
            with open(path, "rb") as f:
                st.download_button("Download .apkg", f, file_name=path)