import streamlit as st
from transformers import pipeline
from fpdf import FPDF
from googletrans import Translator
from gtts import gTTS
import base64
import tempfile

# Load Hugging Face model
@st.cache_resource
def load_model():
    return pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.1")

generator = load_model()

# Translator
translator = Translator()

# Page config
st.set_page_config(page_title="Explain Like I'm 5", page_icon="🧸", layout="centered")
st.markdown("<h1 style='text-align: center;'>🧸 Explain Like I'm 5</h1>", unsafe_allow_html=True)
st.markdown("<p style='text-align: center;'>Ask anything and I’ll explain it super simply 👶</p>", unsafe_allow_html=True)

# Input
user_input = st.text_input("🎯 Enter a topic or question:", placeholder="e.g., What is blockchain?")
language = st.selectbox("🌐 Choose output language:", ["English", "Hindi", "Marathi"])

with st.expander("💡 Try These Examples"):
    st.markdown("- What is AI?\n- Why is the sky blue?\n- How does Wi-Fi work?\n- What is climate change?")

# Hugging Face response
def generate_eli5_response(topic):
    prompt = f"Explain this to a 5-year-old: {topic}"
    result = generator(prompt, max_new_tokens=150, do_sample=True, temperature=0.7)
    return result[0]['generated_text'].replace(prompt, "").strip()

# Translate
def translate_text(text, lang_code):
    return translator.translate(text, dest=lang_code).text

# Language map
lang_map = {
    "English": "en",
    "Hindi": "hi",
    "Marathi": "mr"
}

# PDF Export
def export_to_pdf(topic, explanation):
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=12)
    pdf.multi_cell(0, 10, f"Topic: {topic}\n\nExplanation:\n{explanation}")
    return pdf.output(dest='S').encode('latin-1')

# Text-to-Speech
def text_to_speech(text, lang_code):
    tts = gTTS(text, lang=lang_code)
    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp:
        tts.save(tmp.name)
        audio_path = tmp.name
    return audio_path

# History
if 'history' not in st.session_state:
    st.session_state['history'] = []

# Button logic
if st.button("✨ Explain it to me!"):
    if user_input.strip() == "":
        st.warning("Please enter a topic.")
    else:
        with st.spinner("Explaining like you're 5..."):
            explanation = generate_eli5_response(user_input)

            # Translate if needed
            lang_code = lang_map[language]
            if language != "English":
                explanation_translated = translate_text(explanation, lang_code)
            else:
                explanation_translated = explanation

            # Save to history
            st.session_state['history'].insert(0, {
                "topic": user_input,
                "language": language,
                "explanation": explanation_translated
            })
            st.session_state['history'] = st.session_state['history'][:5]

            # Display result
            st.success("🍼 Here's your explanation:")
            st.markdown(f"**{explanation_translated}**")

            # TTS playback
            audio_path = text_to_speech(explanation_translated, lang_code)
            with open(audio_path, "rb") as audio_file:
                audio_bytes = audio_file.read()
                st.audio(audio_bytes, format="audio/mp3")

            # Export to PDF
            pdf_data = export_to_pdf(user_input, explanation_translated)
            st.download_button("📄 Download as PDF", data=pdf_data, file_name=f"ELI5-{user_input[:30]}.pdf", mime="application/pdf")

# Show history
if st.session_state['history']:
    with st.expander("📜 Past Explanations"):
        for i, entry in enumerate(st.session_state['history']):
            st.markdown(f"**{i+1}. {entry['topic']} ({entry['language']})**")
            st.markdown(f"> {entry['explanation']}")

# Footer
st.markdown("---")
st.markdown("<p style='text-align: center;'>❤️ Made with Love. By Akash Shahade</p>", unsafe_allow_html=True)