import streamlit as st
import os
import tempfile
import uuid
import fitz  # PyMuPDF
import easyocr
import whisper
import docx
import yt_dlp
import csv
import genanki
from transformers import pipeline

# === Helper Functions ===

def process_pdf(path):
    text = ""
    doc = fitz.open(path)
    reader = easyocr.Reader(['en'], gpu=False)
    for page in doc:
        t = page.get_text()
        if t.strip():
            text += t
        else:
            pix = page.get_pixmap()
            img_path = f"/tmp/{uuid.uuid4()}.png"
            pix.save(img_path)
            result = reader.readtext(img_path, detail=0)
            text += "\n".join(result)
    return text

def process_image(path):
    reader = easyocr.Reader(['en'], gpu=False)
    result = reader.readtext(path, detail=0)
    return "\n".join(result)

def process_audio(path):
    model = whisper.load_model("base")
    result = model.transcribe(path)
    return result["text"]

def process_text(path):
    if path.endswith(".txt"):
        with open(path, "r", encoding="utf-8") as f:
            return f.read()
    elif path.endswith(".docx"):
        doc = docx.Document(path)
        return "\n".join([para.text for para in doc.paragraphs])
    return ""

def process_youtube(url):
    temp_dir = tempfile.gettempdir()
    audio_path = os.path.join(temp_dir, f"{uuid.uuid4()}.mp3")
    ydl_opts = {
        'format': 'bestaudio/best',
        'outtmpl': audio_path,
        'postprocessors': [{
            'key': 'FFmpegExtractAudio',
            'preferredcodec': 'mp3',
            'preferredquality': '192',
        }],
        'quiet': True,
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        ydl.download([url])
    return process_audio(audio_path)

@st.cache_resource
def load_llm_swarm():
    return {
        "fast": pipeline("text2text-generation", model="google/flan-t5-small"),
        "bio": pipeline("text2text-generation", model="mrm8488/t5-base-finetuned-question-generation-ap"),
        "deep": pipeline("text2text-generation", model="google/flan-t5-base"),
        "mistral": pipeline("text2text-generation", model="google/flan-t5-large"),
        "fallback": pipeline("text2text-generation", model="MBZUAI/LaMini-Flan-T5-248M")
    }

def generate_flashcards(text, types=["Q&A"], max_cards=100):
    from random import choice
    llm_swarm = load_llm_swarm()
    chunks = [text[i:i + 400] for i in range(0, len(text), 400)][:max_cards]
    prompts, tags = [], []

    for chunk in chunks:
        if "Q&A" in types:
            prompts.append(f"Generate a question and answer:\n{chunk}")
            tags.append("Q&A")
        if "Cloze" in types:
            prompts.append(f"Make a cloze deletion from:\n{chunk}")
            tags.append("Cloze")
        if "MCQ" in types:
            prompts.append(f"Generate a multiple choice question:\n{chunk}")
            tags.append("MCQ")
        if "Reverse" in types:
            prompts.append(f"Generate a question and answer:\n{chunk}")
            tags.append("Reverse")

    cards = []
    for i, prompt in enumerate(prompts):
        engine_name = choice(list(llm_swarm.keys()))
        engine = llm_swarm[engine_name]
        tag = tags[i]
        try:
            output = engine(prompt, max_length=128)[0]["generated_text"]
        except:
            output = llm_swarm["fallback"](prompt, max_length=64)[0]["generated_text"]

        if tag in ["Q&A", "Reverse"]:
            q, a = (output.split(":", 1) + [""])[:2]
            if tag == "Reverse":
                q, a = a.strip(), q.strip()
            cards.append({"question": q.strip(), "answer": a.strip(), "tag": tag})
        elif tag == "Cloze":
            cards.append({"question": output.strip(), "answer": "[...]", "tag": tag})
        elif tag == "MCQ":
            cards.append({"question": output.strip(), "answer": "Choose best option", "tag": tag})

    return cards

def export_to_csv(cards, filename="batanki_cards.csv"):
    with open(filename, "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["Question", "Answer", "Type"])
        for card in cards:
            writer.writerow([card["question"], card["answer"], card["tag"]])

def export_to_apkg(cards, deck_name="BatAnkiDeck"):
    deck_id = int(uuid.uuid4()) >> 64
    model = genanki.Model(
        1607392319,
        "BatAnkiModel",
        fields=[{"name": "Question"}, {"name": "Answer"}],
        templates=[{
            "name": "Card 1",
            "qfmt": "{{Question}}",
            "afmt": "{{FrontSide}}<hr id='answer'>{{Answer}}",
        }]
    )
    deck = genanki.Deck(deck_id, deck_name)
    for card in cards:
        note = genanki.Note(model=model, fields=[card["question"], card["answer"]])
        deck.add_note(note)
    output_path = f"{deck_name}.apkg"
    genanki.Package(deck).write_to_file(output_path)
    return output_path

# === Streamlit UI ===

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 file", type=["pdf", "txt", "docx", "jpg", "png", "mp3", "wav"])
youtube_url = st.sidebar.text_input("Or paste YouTube link")

deck_name = st.text_input("Deck Name", value="BatAnkiDeck")
types_selected = st.multiselect("Flashcard Types", ["Q&A", "Cloze", "MCQ", "Reverse"], default=["Q&A"])
max_cards = st.slider("Max Cards", 5, 500, 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("Select 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("CSV exported.")
    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)