"""
NCLEX Practice Question Generator
Streamlit app — Hugging Face Spaces (free CPU tier)
"""

import random
import streamlit as st

from questions.bank        import filter_questions, get_categories, get_all
from questions.calculations import generate_batch
from scorer                import calculate_score, category_percentage

# ---------------------------------------------------------------------------
# Page config
# ---------------------------------------------------------------------------
st.set_page_config(
    page_title="NCLEX Practice — Student Nurses",
    page_icon="📋",
    layout="wide",
    initial_sidebar_state="expanded",
)

# ---------------------------------------------------------------------------
# CSS
# ---------------------------------------------------------------------------
st.markdown("""
<style>
.question-card {
    background:#f8fafc; border:1px solid #d0dae8;
    border-radius:10px; padding:1.2rem 1.4rem; margin-bottom:1rem;
}
.correct-ans  { background:#e8f5e9; border-left:4px solid #2e7d32;
                padding:0.6rem 1rem; border-radius:4px; margin-top:0.5rem; }
.wrong-ans    { background:#fce4ec; border-left:4px solid #c62828;
                padding:0.6rem 1rem; border-radius:4px; margin-top:0.5rem; }
.rationale    { background:#e3f2fd; border-left:4px solid #1565c0;
                padding:0.6rem 1rem; border-radius:4px; margin-top:0.5rem; }
.sata-badge   { background:#fff8e1; color:#e65100; padding:2px 8px;
                border-radius:4px; font-size:0.75em; font-weight:700; }
.calc-badge   { background:#f3e5f5; color:#6a1b9a; padding:2px 8px;
                border-radius:4px; font-size:0.75em; font-weight:700; }
.score-big    { font-size:3rem; font-weight:800; text-align:center; }
.diff-easy    { color:#2e7d32; font-weight:600; font-size:0.78em; }
.diff-inter   { color:#e65100; font-weight:600; font-size:0.78em; }
.diff-hard    { color:#c62828; font-weight:600; font-size:0.78em; }
</style>
""", unsafe_allow_html=True)

# ---------------------------------------------------------------------------
# Session state
# ---------------------------------------------------------------------------
_DEFAULTS = {
    "quiz_questions":  [],
    "quiz_answers":    {},
    "quiz_submitted":  False,
    "quiz_score":      None,
    "calc_questions":  [],
    "calc_answers":    {},
    "calc_submitted":  False,
    "current_q_idx":   0,
    "quiz_mode":       "all",   # "all" or "one-by-one"
}
for k, v in _DEFAULTS.items():
    if k not in st.session_state:
        st.session_state[k] = v

CATEGORIES = get_categories()
DIFFICULTY_LABELS = {"beginner": "🟢 Beginner", "intermediate": "🟡 Intermediate", "advanced": "🔴 Advanced"}
DIFF_CSS   = {"beginner": "diff-easy", "intermediate": "diff-inter", "advanced": "diff-hard"}


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _diff_badge(diff: str) -> str:
    css   = DIFF_CSS.get(diff, "diff-inter")
    label = DIFFICULTY_LABELS.get(diff, diff.title())
    return f'<span class="{css}">{label}</span>'


def render_question(q: dict, idx: int, key_prefix: str, show_answer: bool = False):
    """Render a question card. Returns the user's answer or None."""
    qtype = q["type"]

    with st.container():
        # Header row
        hcol1, hcol2, hcol3 = st.columns([4, 1, 1])
        with hcol1:
            st.markdown(f"**Question {idx + 1}** — {q['category']} › {q['subcategory']}")
        with hcol2:
            st.markdown(_diff_badge(q.get("difficulty", "")), unsafe_allow_html=True)
        with hcol3:
            if qtype == "sata":
                st.markdown('<span class="sata-badge">SELECT ALL</span>', unsafe_allow_html=True)
            elif qtype == "calculation":
                st.markdown('<span class="calc-badge">CALCULATION</span>', unsafe_allow_html=True)

        # Stem
        st.markdown(f"**{q['stem']}**")

        # Answer input
        user_answer = None
        options     = q["options"]

        if qtype == "sata":
            st.caption("*Select all that apply*")
            selected = []
            for oi, opt in enumerate(options):
                if show_answer:
                    is_correct_opt = oi in q["correct"]
                    icon = "✅" if is_correct_opt else "❌"
                    st.markdown(f"{icon} {opt}")
                else:
                    chk = st.checkbox(opt, key=f"{key_prefix}_q{idx}_o{oi}")
                    if chk:
                        selected.append(oi)
            user_answer = selected

        else:  # mcq or calculation
            if show_answer:
                correct_txt = options[q["correct"]]
                for oi, opt in enumerate(options):
                    if oi == q["correct"]:
                        st.markdown(f"✅ **{opt}**")
                    else:
                        st.markdown(f"○ {opt}")
            else:
                choice = st.radio(
                    "Select your answer:",
                    options,
                    key=f"{key_prefix}_q{idx}",
                    index=None,
                    label_visibility="collapsed",
                )
                user_answer = options.index(choice) if choice else None

        # Show rationale if reviewing
        if show_answer:
            st.markdown(
                f'<div class="rationale">📖 <b>Rationale:</b> {q["rationale"]}</div>',
                unsafe_allow_html=True,
            )

        st.divider()
        return user_answer


# ---------------------------------------------------------------------------
# Sidebar
# ---------------------------------------------------------------------------
with st.sidebar:
    st.markdown("## 📋 NCLEX Practice")
    st.divider()

    st.markdown("**Quiz Settings**")
    sel_cats = st.multiselect(
        "Categories",
        CATEGORIES,
        default=CATEGORIES,
        help="Select which NCLEX categories to include",
    )
    sel_diff = st.selectbox(
        "Difficulty",
        ["Any", "beginner", "intermediate", "advanced"],
        format_func=lambda x: "Any difficulty" if x == "Any" else DIFFICULTY_LABELS[x],
    )
    sel_type = st.selectbox(
        "Question type",
        ["Any", "mcq", "sata", "calculation"],
        format_func=lambda x: {
            "Any": "All types", "mcq": "Multiple Choice (MCQ)",
            "sata": "Select All That Apply (SATA)", "calculation": "Dosage Calculation"
        }[x],
    )
    n_questions = st.slider("Number of questions", 5, 30, 10)

    st.divider()
    available = filter_questions(sel_cats, sel_diff, sel_type if sel_type != "calculation" else None)
    st.caption(f"📚 {len(get_all())} questions in bank · {len(available)} match filters")

    st.divider()
    st.markdown(
        """
**NCLEX-RN Test Plan**
- Safe & Effective Care Environment
- Health Promotion & Maintenance
- Psychosocial Integrity
- Physiological Integrity
        """
    )


# ---------------------------------------------------------------------------
# Header
# ---------------------------------------------------------------------------
st.title("📋 NCLEX Practice Question Generator")
st.caption(
    "Evidence-based NCLEX-style practice questions · MCQ · Select All That Apply · "
    "Dosage Calculations · Full rationales for every question"
)

# ---------------------------------------------------------------------------
# Tabs
# ---------------------------------------------------------------------------
tab_quiz, tab_calc, tab_results, tab_review = st.tabs(
    ["🎯 Practice Quiz", "💉 Dosage Calculations", "📊 My Results", "🔁 Review Wrong Answers"]
)

# ========================= PRACTICE QUIZ TAB ==============================
with tab_quiz:
    col_start, col_reset = st.columns([3, 1])

    with col_start:
        if st.button("🎯 Generate New Quiz", type="primary", use_container_width=True):
            bank = filter_questions(sel_cats, sel_diff, sel_type if sel_type != "calculation" else None)
            if not bank:
                st.error("No questions match your filters. Adjust the sidebar settings.")
            else:
                sampled = random.sample(bank, min(n_questions, len(bank)))
                # Mix in calc questions if type is Any or calculation
                if sel_type in ("Any", "calculation"):
                    n_calc = max(1, n_questions // 5)
                    calc_q = generate_batch(n_calc)
                    if sel_type == "calculation":
                        sampled = calc_q * ((n_questions // n_calc) + 1)
                        sampled = sampled[:n_questions]
                    else:
                        sampled = sampled[:-n_calc] + calc_q
                        random.shuffle(sampled)

                st.session_state.quiz_questions  = sampled[:n_questions]
                st.session_state.quiz_answers    = {}
                st.session_state.quiz_submitted  = False
                st.session_state.quiz_score      = None
                st.rerun()

    with col_reset:
        if st.button("🔄 Reset", use_container_width=True):
            for k in ["quiz_questions","quiz_answers","quiz_submitted","quiz_score"]:
                st.session_state[k] = [] if k == "quiz_questions" else ({} if "answers" in k else (False if "submitted" in k else None))
            st.rerun()

    questions = st.session_state.quiz_questions

    if not questions:
        st.info("Click **Generate New Quiz** to start practising.")
    else:
        submitted = st.session_state.quiz_submitted

        if not submitted:
            st.markdown(f"**{len(questions)} questions** — read each carefully, then click Submit.")
            st.divider()

            # Render all questions
            for i, q in enumerate(questions):
                ans = render_question(q, i, "quiz", show_answer=False)
                if ans is not None and ans != []:
                    st.session_state.quiz_answers[i] = ans

            if st.button("✅ Submit Quiz", type="primary", use_container_width=False):
                score = calculate_score(st.session_state.quiz_answers, questions)
                st.session_state.quiz_score     = score
                st.session_state.quiz_submitted = True
                st.rerun()

        else:
            # Show results summary
            score = st.session_state.quiz_score
            pct   = score["percentage"]
            band  = score["performance_band"]
            emoji = "🏆" if pct >= 75 else "📚"

            c1, c2, c3 = st.columns(3)
            c1.metric("Score", f"{score['correct']} / {score['total']}")
            c2.metric("Percentage", f"{pct}%")
            c3.metric("Result", band)

            st.progress(pct / 100)
            st.info(f"{emoji} {score['feedback']}")
            st.divider()

            # Show all questions with correct answers
            st.subheader("Question Review")
            for i, q in enumerate(questions):
                user_ans = st.session_state.quiz_answers.get(i)
                correct  = q["correct"]
                is_right = (set(user_ans) == set(correct)) if q["type"] == "sata" else user_ans == correct

                icon = "✅" if is_right else "❌"
                with st.expander(f"{icon} Q{i+1}: {q['stem'][:80]}…", expanded=not is_right):
                    render_question(q, i, "review", show_answer=True)


# ========================= DOSAGE CALCULATIONS TAB ========================
with tab_calc:
    c_gen, c_rst = st.columns([3, 1])
    with c_gen:
        if st.button("💉 Generate Calculation Set", type="primary", use_container_width=True):
            st.session_state.calc_questions = generate_batch(n_questions)
            st.session_state.calc_answers   = {}
            st.session_state.calc_submitted = False
            st.rerun()
    with c_rst:
        if st.button("🔄 Reset ", use_container_width=True, key="calc_rst"):
            st.session_state.calc_questions = []
            st.session_state.calc_answers   = {}
            st.session_state.calc_submitted = False
            st.rerun()

    calc_qs = st.session_state.calc_questions

    if not calc_qs:
        st.info(
            "Click **Generate Calculation Set** for a fresh set of dosage math problems. "
            "Each set is uniquely generated — unlimited practice!"
        )
        with st.expander("📐 Dosage Calculation Formula Reference"):
            st.markdown("""
**Oral Tablets:**
> (Ordered dose ÷ Dose on hand) × Quantity = Tablets to give

**IV Rate (mL/hr):**
> Volume (mL) ÷ Time (hours) = mL/hr

**Manual Drip Rate (gtt/min):**
> (Volume × Drop factor) ÷ Time (minutes) = gtt/min

**Weight-based IV Infusion:**
> (mcg/kg/min × kg × 60) ÷ Concentration (mcg/mL) = mL/hr

**Paediatric Oral Dose:**
> mg/kg × weight (kg) ÷ concentration (mg/mL) = mL to give
            """)
    else:
        calc_submitted = st.session_state.calc_submitted

        if not calc_submitted:
            for i, q in enumerate(calc_qs):
                ans = render_question(q, i, "calc", show_answer=False)
                if ans is not None and ans != []:
                    st.session_state.calc_answers[i] = ans

            if st.button("✅ Submit Calculations", type="primary"):
                score = calculate_score(st.session_state.calc_answers, calc_qs)
                st.session_state.calc_submitted = True
                st.session_state.quiz_score     = score
                st.rerun()
        else:
            score = calculate_score(st.session_state.calc_answers, calc_qs)
            pct   = score["percentage"]
            c1, c2, c3 = st.columns(3)
            c1.metric("Correct", f"{score['correct']} / {score['total']}")
            c2.metric("Percentage", f"{pct}%")
            c3.metric("Band", score["performance_band"])
            st.progress(pct / 100)
            st.divider()

            for i, q in enumerate(calc_qs):
                user_ans = st.session_state.calc_answers.get(i)
                is_right = user_ans == q["correct"]
                icon = "✅" if is_right else "❌"
                with st.expander(f"{icon} Q{i+1}: {q['stem'][:80]}…", expanded=not is_right):
                    render_question(q, i, "calc_rev", show_answer=True)


# ========================= RESULTS TAB ====================================
with tab_results:
    score = st.session_state.quiz_score
    if not score:
        st.info("Complete a quiz to see your results here.")
    else:
        pct  = score["percentage"]
        band = score["performance_band"]

        st.markdown(f'<div class="score-big">{pct}%</div>', unsafe_allow_html=True)
        st.markdown(f"<h3 style='text-align:center'>{band}</h3>", unsafe_allow_html=True)
        st.markdown(f"<p style='text-align:center'>{score['feedback']}</p>", unsafe_allow_html=True)
        st.divider()

        # Category breakdown
        st.subheader("Performance by Category")
        cat_pct = category_percentage(score)
        for cat, p in sorted(cat_pct.items(), key=lambda x: x[1]):
            col_l, col_r = st.columns([3, 1])
            col_l.write(cat)
            col_r.write(f"**{p}%**")
            st.progress(p / 100)

        st.divider()

        # Difficulty breakdown
        st.subheader("Performance by Difficulty")
        diff_data = score["by_difficulty"]
        dcols = st.columns(3)
        for i, (diff, data) in enumerate(diff_data.items()):
            t = data["total"]
            p = round((data["correct"] / t) * 100, 1) if t > 0 else 0
            dcols[i % 3].metric(
                DIFFICULTY_LABELS.get(diff, diff.title()),
                f"{p}%",
                f"{data['correct']}/{t} correct",
            )


# ========================= REVIEW TAB =====================================
with tab_review:
    score = st.session_state.quiz_score
    wrong = score["wrong_questions"] if score else []

    if not wrong:
        msg = "No wrong answers to review — great work! 🎉" if score else "Complete a quiz first."
        st.info(msg)
    else:
        st.subheader(f"Review: {len(wrong)} question(s) to revisit")
        st.caption("Study the rationale for each — understanding WHY is the key to NCLEX success.")
        st.divider()

        for i, item in enumerate(wrong):
            q = item["question"]
            with st.expander(f"❌ {q['category']} | {q['stem'][:70]}…"):
                st.markdown(f"**Category:** {q['category']} › {q['subcategory']}")
                st.markdown(f"**NCLEX Framework:** {q.get('nclex_framework','')}")
                st.markdown(_diff_badge(q.get("difficulty", "")), unsafe_allow_html=True)
                st.divider()
                render_question(q, i, "wrong_rev", show_answer=True)

# ---------------------------------------------------------------------------
# Footer
# ---------------------------------------------------------------------------
st.divider()
st.caption(
    "Questions aligned to the 2023 NCLEX-RN Test Plan (NCSBN). "
    "Dosage calculations are dynamically generated — every set is unique. "
    "For educational purposes only — always apply clinical judgment in practice."
)