app.py

import streamlit as st
import tempfile
from main import run_question_pipeline
from utils import extract_text_from_pdf, extract_text_from_docx, generate_pdf
from markdown import markdown
import re
from docx import Document


def strip_markdown(md_text):
    html = markdown(md_text)
    clean_text = re.sub('<[^<]+?>', '', html)  # remove HTML tags
    return clean_text

st.set_page_config(page_title="AI Question Paper Generator", layout="centered")
st.title("📄 AI Question Paper Generator")

uploaded_file = st.file_uploader("Upload Study Material (PDF or DOCX)", type=["pdf", "docx"])

if uploaded_file:
    with tempfile.NamedTemporaryFile(delete=False, suffix=uploaded_file.name[-4:]) as tmp_file:
        tmp_file.write(uploaded_file.read())
        file_path = tmp_file.name

    # Extract text from file
    if uploaded_file.name.endswith(".pdf"):
        raw_text = extract_text_from_pdf(file_path)
    else:
        raw_text = extract_text_from_docx(file_path)

    st.success("📚 Study material loaded successfully!")

    col1, col2 = st.columns([2, 3])
    with col1:
        st.markdown(
            "<h4 style='font-size:1.66rem;'>🎯 Select Difficulty </h4>",
            unsafe_allow_html=True
        )

    st.markdown("""
    <style>
    div[data-baseweb="select"] {
        width: 65% !important;
        float: right;
    }
    </style>
    """, unsafe_allow_html=True)

    with col2:
        difficulty = st.selectbox(
            label=" ",  # Empty label
            options=["Easy", "Medium", "Hard"],
            index=1,
            key="difficulty_level"
        )

    col1, col2 = st.columns([2, 3])
    with col1:
        st.markdown(
            "<h4 style='font-size:1.66rem;'>🔢 Select Total Marks</h4>",
            unsafe_allow_html=True
        )

    st.markdown("""
    <style>
    div[data-baseweb="select"] {
        width: 65% !important;
        float: right;
    }
    </style>
    """, unsafe_allow_html=True)

    with col2:
        total_marks = st.selectbox(
            label=" ",
            options=[10, 20, 30, 40, 50, 60, 70, 80, 90, 100],
            index=0,
            key="total_marks"
        )

    st.subheader("✍️ Customize Question Types")

    # Row 1: MCQ, 1-mark, 2-mark
    col1, col2, col3 = st.columns(3)
    with col1:
        mcq_count = st.number_input("MCQs (1 mark)", min_value=0, value=0)
    with col2:
        q1 = st.number_input("1-mark Questions", min_value=0, value=0)
    with col3:
        q2 = st.number_input("2-mark Questions", min_value=0, value=0)

    # Row 2: 3-mark, 4-mark, 5-mark
    col4, col5, col6 = st.columns(3)
    with col4:
        q3 = st.number_input("3-mark Questions", min_value=0, value=0)
    with col5:
        q4 = st.number_input("4-mark Questions", min_value=0, value=0)
    with col6:
        q5 = st.number_input("5-mark Questions", min_value=0, value=0)

    # Row 3: 10-mark, 15-mark, 20-mark
    col7, col8, col9 = st.columns(3)
    with col7:
        q10 = st.number_input("10-mark Questions", min_value=0, value=0)
    with col8:
        q15 = st.number_input("15-mark Questions", min_value=0, value=0)
    with col9:
        q20 = st.number_input("20-mark Questions", min_value=0, value=0)

    actual_total = (
        mcq_count * 1 + q1 * 1 + q2 * 2 + q3 * 3 + q4 * 4 +
        q5 * 5 + q10 * 10 + q15 * 15 + q20 * 20
    )

    if actual_total != total_marks:
        st.warning(f"⚠️ Total marks ({actual_total}) do not match selected value ({total_marks}). Please adjust.")

    if st.button("🚀 Generate Question Paper"):
        with st.spinner("Generating... Please wait..."):
            counts = {
                "mcq": mcq_count,
                "q1": q1,
                "q2": q2,
                "q3": q3,
                "q4": q4,
                "q5": q5,
                "q10": q10,
                "q15": q15,
                "q20": q20
            }
            result = run_question_pipeline(raw_text, total_marks, counts, difficulty)

        if "error" in result:
            st.error(f"❌ Error: {result['error']}")
        else:
            st.subheader("📘 Generated Question Paper")
            st.markdown(result["questions"])


            clean_text = strip_markdown(result["questions"])
            pdf_bytes = generate_pdf("not_used", clean_text)

            st.download_button(
                label="📥 Download Question Paper (PDF)",
                data=pdf_bytes,
                file_name="question_paper.pdf",
                mime="application/pdf"
            )