import streamlit as st
import json

# Import your Orchestrator and PDF Exporter
from pipeline import QuizOrchestrator
import PDF_Exporter8 as exporter

# ==========================================
# PAGE CONFIGURATION
# ==========================================
st.set_page_config(
    page_title="Neural Quiz Platform", 
    page_icon="🧠", 
    layout="wide",
    initial_sidebar_state="expanded"
)

# ==========================================
# INITIALIZE BACKEND (Cached to save RAM)
# ==========================================
@st.cache_resource
def load_orchestrator():
    return QuizOrchestrator()

# This ensures the heavy models only load once when the app starts
orchestrator = load_orchestrator()

# ==========================================
# SIDEBAR: SETTINGS & INPUTS
# ==========================================
st.sidebar.title("⚙️ Quiz Configuration")

# 1. Input Mode
input_mode = st.sidebar.radio(
    "Select Input Mode:", 
    ["Topic Search (Wikipedia)", "Custom Paragraph"]
)

topic_input = ""
custom_text = ""

if input_mode == "Topic Search (Wikipedia)":
    topic_input = st.sidebar.text_input(
        "Enter Topic Name (Required):", 
        placeholder="e.g., Machine Learning",
        help="The system will search your local cache or crawl Wikipedia."
    )
else:
    topic_input = st.sidebar.text_input(
        "Enter a Label for this Topic:", 
        placeholder="e.g., Chapter 1 Biology",
        help="Give this text a name so we can save it to your local database."
    )
    custom_text = st.sidebar.text_area(
        "Paste your paragraphs here:", 
        height=250,
        help="Paste the exact text you want the AI to read."
    )

st.sidebar.markdown("---")

# 2. Quiz Parameters
num_questions = st.sidebar.slider("Number of Questions", min_value=1, max_value=10, value=3)
q_types = st.sidebar.multiselect(
    "Question Types", 
    ["MCQ", "FIB", "TF"], 
    default=["MCQ", "FIB", "TF"]
)

# ==========================================
# MAIN WINDOW
# ==========================================
st.title("🧠 Neural Assessment Generator")
st.markdown("""
Transform any topic or text into a professional, multi-format educational assessment. 
Powered by T5, RoBERTa, Sense2Vec, and Sentence-Transformers.
""")

# Generation Button
if st.button("🚀 Generate Quiz", type="primary", use_container_width=True):
    
    # Input Validation
    if not topic_input:
        st.error("⚠️ Please provide a Topic Name to proceed.")
    elif input_mode == "Custom Paragraph" and not custom_text.strip():
        st.error("⚠️ Please paste some text into the custom paragraph box.")
    elif not q_types:
        st.error("⚠️ Please select at least one Question Type.")
    else:
        # Show a loading spinner while the backend works
        with st.spinner(f"Compiling intelligence for '{topic_input}'. This requires heavy neural processing and may take a moment..."):
            
            # Call the Master Orchestrator
            generated_quiz = orchestrator.create_quiz(
                topic_name=topic_input,
                custom_paragraphs=custom_text if input_mode == "Custom Paragraph" else None,
                num_questions=num_questions,
                question_types=q_types
            )
            
            if not generated_quiz:
                st.error("❌ Failed to generate questions. The content might be too short, or the AI rejected the generated questions for low quality.")
            else:
                st.success(f"✅ Successfully generated {len(generated_quiz)} questions!")
                
                # ==========================================
                # PDF EXPORT BUTTON
                # ==========================================
                # Generate the PDF in memory
                pdf_bytes = exporter.generate_pdf(topic_input, generated_quiz)
                
                # Create the download button
                st.download_button(
                    label="📄 Download Quiz as PDF (with Answer Key)",
                    data=pdf_bytes,
                    file_name=f"{topic_input.replace(' ', '_')}_Assessment.pdf",
                    mime="application/pdf",
                    type="secondary"
                )
                
                st.markdown("---")
                
                # ==========================================
                # DISPLAY THE QUIZ INTERACTIVELY
                # ==========================================
                st.subheader("Interactive Preview")
                
                for i, q in enumerate(generated_quiz):
                    # Use an expander for each question for a clean UI
                    with st.expander(f"Question {i+1} — [{q['type']}]", expanded=True):
                        
                        if q['type'] == "MCQ":
                            st.markdown(f"**{q['question']}**")
                            for idx, opt in enumerate(q['options']):
                                # Highlight the correct answer visually on the screen
                                if opt == q['answer']:
                                    st.markdown(f"- ✅ **{opt}** *(Correct Answer)*")
                                else:
                                    st.markdown(f"- ⚪ {opt}")
                                    
                        elif q['type'] == "FIB":
                            st.markdown(f"**Fill in the blank:**")
                            st.info(f"{q['question']}")
                            st.markdown(f"**Answer:** ✅ {q['answer']}")
                            
                        elif q['type'] == "TF":
                            st.markdown(f"**True or False?**")
                            st.warning(f"{q['statement']}")
                            st.markdown(f"**Answer:** ✅ {q['answer']}")