# phase/Student_view/mini_quiz.py
import datetime
from typing import List, Dict, Any, Optional
import streamlit as st

# Reuse the unified backend client
from utils import api as backend_api

# ---------------------------------------------
# Public quiz-state helpers
# ---------------------------------------------
QUIZ_SS_DEFAULTS = {
    "quiz_data": None,           # original quiz payload (list[dict])
    "quiz_answers": {},          # q_index -> "A"|"B"|"C"|"D"
    "quiz_result": None,         # backend result dict
    "chatbot_feedback": None,    # str
    "_auto_quiz_started": False, # helper to avoid double-start
}


def ensure_quiz_state():
    """Ensure quiz-related keys exist in st.session_state."""
    for k, v in QUIZ_SS_DEFAULTS.items():
        if k not in st.session_state:
            st.session_state[k] = v


# ---------------------------------------------
# Backend calls (isolated here)
# ---------------------------------------------

# --- in start_quiz() ---
def start_quiz(level: str, module_id: int, lesson_title: str) -> bool:
    ensure_quiz_state()
    try:
        resp = backend_api.generate_quiz(
            lesson_id=module_id,
            level_slug=level,
            lesson_title=lesson_title,
        )
    except Exception as e:
        st.error(f"Could not generate quiz: {e}")
        return False

    # Accept both the old (list) and new (dict with 'items') shapes
    items = resp.get("items") if isinstance(resp, dict) else resp
    if not isinstance(items, list) or not items:
        st.error("Quiz could not be generated. Please try again.")
        return False

    # Normalize: ensure id + answer_key are present
    normalized = []
    for i, q in enumerate(items, start=1):
        qid = q.get("id") or q.get("qid") or q.get("position") or f"q{i}"
        answer_key = q.get("answer_key") or q.get("answer")  # backend may send either
        normalized.append({
            "id": qid,
            "question": (q.get("question") or "").strip(),
            "options": q.get("options") or [],
            "answer_key": answer_key,  # keep for grading payload
        })

    st.session_state.quiz_data = normalized
    st.session_state.quiz_answers = {}
    st.session_state.mode = "quiz"
    return True


def submit_quiz(level: str, module_id: int, original_quiz: List[Dict[str, Any]], answers_map: Dict[int, str]) -> Optional[Dict[str, Any]]:
    """Submit answers and return the grading result dict."""
    # --- in submit_quiz() ---
    user_answers = []
    for q in original_quiz:
        qid = q.get("id") or q.get("question")  # fallback
        user_answers.append({"id": qid, "answer": answers_map.get(original_quiz.index(q), "")})


    # student_id is required by utils.api.submit_quiz
    student_id = int(((st.session_state.get("user") or {}).get("user_id") or 0))

    try:
        result = backend_api.submit_quiz(
            student_id=student_id,
            lesson_id=module_id,
            level_slug=level,
            user_answers=user_answers,
            original_quiz=original_quiz,
        )
        return result
    except Exception as e:
        st.error(f"Could not submit quiz: {e}")
        return None


def send_quiz_summary_to_chatbot(level: str, module_id: int, lesson_title: str, result: Dict[str, Any]):
    """Send a concise summary to the chatbot and navigate there."""
    score = result.get("score", {})
    correct = int(score.get("correct", 0))
    total = int(score.get("total", 0))
    feedback = (result.get("feedback") or st.session_state.get("chatbot_feedback") or "").strip()

    user_prompt = (
        f"I just finished the quiz for '{lesson_title}' (module {module_id}) "
        f"and scored {correct}/{total}. Please give me 2–3 targeted tips and 1 tiny action "
        f"to improve before the next lesson. If there were wrong answers, explain them simply.\n\n"
        f"Context from grader:\n{feedback}"
    )

    try:
        bot_reply = (backend_api.chat_ai(
            query=user_prompt,
            lesson_id=module_id,
            level_slug=level,
            history=[],
        ) or "").strip()
    except Exception:
        bot_reply = f"(Chatbot unavailable) Based on your result: {feedback or 'Nice work!'}"

    # Seed Chatbot page
    msgs = st.session_state.get("messages") or [{
        "id": "1",
        "text": "Hi! I'm your AI Financial Tutor. What would you like to learn today?",
        "sender": "assistant",
        "timestamp": datetime.datetime.now(),
    }]
    msgs.append({"text": user_prompt, "sender": "user", "timestamp": datetime.datetime.now()})
    msgs.append({"text": bot_reply,   "sender": "assistant", "timestamp": datetime.datetime.now()})

    st.session_state.messages = msgs
    st.session_state.current_page = "Chatbot"


# ---------------------------------------------
# UI helpers
# ---------------------------------------------

def _letter_for(i: int) -> str:
    return chr(ord("A") + i)


def render_quiz(lesson_title: Optional[str] = None):
    """Render the quiz UI (single page). Expects st.session_state.level/module_id to be set."""
    ensure_quiz_state()

    quiz: List[Dict[str, Any]] = st.session_state.quiz_data or []
    if not quiz:
        # No staged quiz — bounce back to lesson
        st.session_state.mode = "lesson"
        st.rerun()

    st.markdown("### Lesson Quiz")

    # Render each question
    for q_idx, q in enumerate(quiz):
        st.markdown(f"**Q{q_idx+1}. {q.get('question','').strip()}**")
        opts = q.get("options") or []

        def _on_select():
            sel = st.session_state[f"ans_{q_idx}"]  # e.g. "A. option text"
            letter = sel.split(".", 1)[0] if isinstance(sel, str) else ""
            st.session_state.quiz_answers[q_idx] = letter

        labels = [f"{_letter_for(i)}. {opt}" for i, opt in enumerate(opts)]
        saved_letter = st.session_state.quiz_answers.get(q_idx)
        pre_idx = next((i for i, l in enumerate(labels) if saved_letter and l.startswith(f"{saved_letter}.")), None)

        st.radio(
            "",
            labels,
            index=pre_idx,
            key=f"ans_{q_idx}",
            on_change=_on_select,
        )
        st.divider()

    all_answered = len(st.session_state.quiz_answers) == len(quiz)
    if st.button("Submit Quiz", disabled=not all_answered):
        with st.spinner("Grading…"):
            result = submit_quiz(
                st.session_state.level,
                st.session_state.module_id,
                quiz,
                st.session_state.quiz_answers,
            )
        if result:
            st.session_state.quiz_result = result
            st.session_state.chatbot_feedback = result.get("feedback")
            # Keep current behavior: jump to Chatbot after grading
            send_quiz_summary_to_chatbot(
                st.session_state.level,
                st.session_state.module_id,
                lesson_title or "This Lesson",
                result,
            )
            st.rerun()


def render_results(planned_topics: Optional[List[str]] = None):
    """Optional results screen (not used by default flow which jumps to Chatbot)."""
    ensure_quiz_state()

    result = st.session_state.quiz_result or {}
    score = result.get("score", {})
    correct = score.get("correct", 0)
    total = score.get("total", 0)

    st.success(f"Quiz Complete! You scored {correct} / {total}.")

    wrong = result.get("wrong", [])
    if wrong:
        with st.expander("Review your answers"):
            for w in wrong:
                st.markdown(f"**{w.get('question','')}**")
                st.write(f"Your answer: {w.get('your_answer','')}")
                st.write(f"Correct answer: {w.get('correct_answer','')}")
                st.divider()

    fb = st.session_state.chatbot_feedback
    if fb:
        st.markdown("#### Tutor Explanation")
        st.write(fb)

    # Navigation controls
    planned_topics = planned_topics or []
    try:
        quiz_index = [t.strip().lower() for t in planned_topics].index("quiz")
    except ValueError:
        quiz_index = None

    c1, c2, c3 = st.columns([1, 1, 1])
    with c1:
        if st.button("Back to Modules"):
            st.session_state.mode = "catalog"
            st.session_state.module_id = None
            st.rerun()
    with c2:
        if st.button("Ask the Chatbot →"):
            st.session_state.current_page = "Chatbot"
            st.session_state.chatbot_prefill = fb
            st.rerun()
    with c3:
        if quiz_index is not None and quiz_index + 1 < len(planned_topics):
            if st.button("Continue Lesson →"):
                st.session_state.mode = "lesson"
                st.session_state.topic_idx = quiz_index + 1
                st.rerun()