#!/usr/bin/env python3
"""
AI Teacher Bot - Single Panel UI (fixed example-evaluation & progression)
"""

import gradio as gr
import io
from contextlib import redirect_stdout

from main import LEVELS, check_api_key, generate_curriculum
from user_state import UserState
from agents.level_assess import LevelAssessmentAgent
from agents.teacher import TeacherAgent
from agents.bloom_assess import BloomsAssessmentAgent
import json

# Global session
current_session = {
    "user": None,
    "curriculum": None,
    "chapter_idx": 0,
    "module_idx": 0,
    "mode": None,  # modes: None/idle, "assessment", "teaching", "awaiting_example", "module_passed", "bloom", "done"
    "questions": [],
    "answers": [],
    "q_idx": 0,
    "bloom_level": None,
    # assessment flow controls
    "correct_count": 0,
    "assessment_feedback": [],  # list of dicts per question with correctness and feedback
    "last_feedback": None,  # stores last question, answer, evaluation, reasoning for challenging
    "challenge_chat_history": [],  # list of [user_msg, assistant_msg] pairs for challenge discussion
    "challenge_exchanges": 0,  # count of challenge exchanges (max 3)
    "challenge_mode": False,  # whether challenge chat is active
    "last_output": ""
}

BLOOM_ORDER = ["remember", "understand", "apply", "analyze", "evaluate", "create"]


def update_main_output(text):
    current_session["last_output"] = text
    return text

# ------------------------------
# Chatbot helpers (Gradio 6 safe)
# ------------------------------
def empty_chat():
    """
    Returns empty chat in messages format (dictionaries with 'role' and 'content' keys).
    """
    return [{"role": "assistant", "content": " "}]


def safe_chat(chat):
    """
    Ensures chat history is always valid. Uses messages format (list of dicts with 'role' and 'content').
    """
    if not isinstance(chat, list):
        return empty_chat()
    if len(chat) == 0:
        return empty_chat()
    # Ensure all items are dicts with role and content
    result = []
    for item in chat:
        if isinstance(item, dict) and "role" in item and "content" in item:
            result.append(item)
        elif isinstance(item, tuple) and len(item) == 2:
            # Convert tuple (user_msg, assistant_msg) to dict format
            result.append({"role": "user", "content": item[0]})
            result.append({"role": "assistant", "content": item[1]})
    return result if result else empty_chat()


def get_output_update():
    return gr.update(value=current_session.get("last_output", ""))

# ------------------------------
# Session & Flow helpers
# ------------------------------
def reset_session_state():
    current_session.update({
        "user": None,
        "curriculum": None,
        "chapter_idx": 0,
        "module_idx": 0,
        "mode": None,
        "questions": [],
        "answers": [],
        "q_idx": 0,
        "bloom_level": None,
        "correct_count": 0,
        "assessment_feedback": [],
        "last_feedback": None,
        "challenge_chat_history": [],
        "challenge_exchanges": 0,
        "challenge_mode": False,
        "last_output": ""
    })

def start_learning_session(topic, claimed_level):
    if not topic or not topic.strip():
        return "❌ Please enter a topic"
    if not check_api_key():
        return "❌ OpenAI API key not configured! Please set OPENAI_API_KEY in your .env file."

    try:
        reset_session_state()
        user = UserState(topic=topic.strip(), claimed_level=claimed_level)
        current_session["user"] = user

        # Generate curriculum (level may be updated later after assessment)
        curriculum = generate_curriculum(user.topic, claimed_level)
        if curriculum is None:
            return "❌ Failed to generate curriculum. Please try again or check your API key."
        current_session["curriculum"] = curriculum
    except Exception as e:
        return f"❌ Error starting session: {str(e)}. Please try again."

    if claimed_level == "novice":
        user.set_actual_level("novice")
        current_session["mode"] = None  # ready to start teaching
        return update_main_output(show_curriculum(curriculum) + "\n\nType 'next' to start Module 1.")
    else:
        # start level assessment
        assessor = LevelAssessmentAgent()
        q_text = assessor.generate_questions(user.topic, claimed_level)
        questions = [q.strip() for q in q_text.split("\n") if q.strip() and q[0].isdigit()]
        if not questions:
            # fallback: skip assessment
            user.set_actual_level(claimed_level)
            current_session["mode"] = None
            return update_main_output(show_curriculum(curriculum) + "\n\nType 'next' to start Module 1.")
        current_session.update({
            "mode": "assessment",
            "questions": questions,
            "q_idx": 0,
            "answers": [],
            "correct_count": 0,
            "assessment_feedback": [],
            "last_feedback": None,  # Will be set after first answer submission
            "challenge_chat_history": [],
            "challenge_exchanges": 0,
            "challenge_mode": False
        })
        # Button should be hidden initially (no feedback yet), will show after first answer
        return update_main_output(f"📝 LEVEL ASSESSMENT\n\nQuestion 1 of {len(questions)}:\n{questions[0]}\n\nPlease submit your answer.")

# ------------------------------
# Level assessment handlers
# ------------------------------
def _strip_number_prefix(q_line: str) -> str:
    # Converts "1. Question" -> "Question" safely
    q = q_line.strip()
    if ". " in q:
        parts = q.split(". ", 1)
        if parts[0].isdigit():
            return parts[1]
    return q


def handle_assessment(answer):
    qs = current_session["questions"]
    idx = current_session["q_idx"]

    if not qs:
        return "⚠️ No assessment in progress."

    question_full = qs[idx]
    question = _strip_number_prefix(question_full)

    # guard on empty/very short answers
    user_answer = (answer or "").strip()
    if len(user_answer.split()) < 5:
        # Don't set last_feedback for invalid answers, so button won't show
        return "❌ Your answer is too brief. Please provide a more detailed and specific response (at least 5 words or 1-2 sentences)."

    assessor = LevelAssessmentAgent()
    try:
        raw = assessor.evaluate_answer(
            current_session["user"].topic,
            current_session["user"].claimed_level,
            question,
            user_answer,
        )
        import json as _json
        parsed = _json.loads(raw)
        evaluation = str(parsed.get("evaluation", "incorrect")).lower()
        reasoning = parsed.get("reasoning", "")
        hint = parsed.get("hint", "")  # Extract hint for completeness
    except Exception as e:
        # Better error handling
        evaluation = "correct" if len(user_answer.split()) >= 20 else "incorrect"
        reasoning = f"Heuristic grading fallback used. (Error: {str(e)})"
        hint = ""

    # Record attempt
    current_session["answers"].append(user_answer)
    is_correct = evaluation == "correct"
    if is_correct:
        current_session["correct_count"] += 1

    # Save per-question feedback
    current_session["assessment_feedback"].append({
        "question": question,
        "correct": is_correct,
        "reason": reasoning,
        "hint": hint,
    })
    
    # Store last feedback for challenging
    current_session["last_feedback"] = {
        "question": question,
        "answer": user_answer,
        "evaluation": evaluation,
        "reasoning": reasoning,
        "is_correct": is_correct,
        "hint": hint
    }

    # Advance to next question or finish
    if idx + 1 < len(qs):
        current_session["q_idx"] += 1
        status = "✅ Correct!" if is_correct else "❌ Incorrect."
        extra = f"\nReason: {reasoning}" if reasoning else ""
        hint_text = f"\n💡 Hint: {hint}" if hint and not is_correct else ""
        return (
            f"{status}{extra}{hint_text}\n\n"
            f"Question {current_session['q_idx']+1} of {len(qs)}:\n{qs[current_session['q_idx']]}"
        )
    else:
        # Finish: compute score and assign level based on claimed level thresholds
        total = len(qs)
        score_pct = (current_session["correct_count"] / total) * 100
        claimed = current_session["user"].claimed_level

        if claimed == "advanced":
            if score_pct >= 70:
                assigned = "advanced"
            elif score_pct >= 65:
                assigned = "intermediate"
            else:
                assigned = "novice"
        elif claimed == "intermediate":
            assigned = "intermediate" if score_pct >= 65 else "novice"
        else:
            assigned = "novice"

        # Build full feedback summary
        lines = [
            "🧪 Assessment Feedback:",
        ]
        for i, fb in enumerate(current_session["assessment_feedback"], 1):
            tag = "✅" if fb["correct"] else "❌"
            line = f"{tag} Q{i}: {fb['question']}"
            if fb.get("reason"):
                line += f"\n   Reason: {fb['reason']}"
            lines.append(line)
        lines.append("")
        lines.append(f"Score: {score_pct:.1f}%  |  Assigned Level: {assigned}")

        current_session["user"].set_actual_level(assigned)
        # Prepare curriculum but show it on next screen
        curriculum = generate_curriculum(current_session["user"].topic, assigned)
        if curriculum:
            current_session["curriculum"] = curriculum

        feedback_text = "\n".join(lines)
        current_session["mode"] = "assessment_summary"
        current_session["last_feedback"] = None  # Clear last feedback when assessment completes
        return feedback_text + "\n\n➡️ Press 'Next' to view your personalized curriculum."

# ------------------------------
# Teaching helpers
# ------------------------------
def show_curriculum(curriculum):
    txt = f"📖 CURRICULUM FOR {current_session['user'].topic.upper()}\n" + "="*40 + "\n"
    for i, ch in enumerate(curriculum.chapters, 1):
        txt += f"\nChapter {i}: {ch.name}\n"
        for j, mod in enumerate(ch.modules, 1):
            txt += f"  {i}.{j} {mod.name}\n"
            if getattr(mod, "learning_objective", None):
                txt += f"     → {mod.learning_objective}\n"
    return txt

def next_step(_):
    mode = current_session["mode"]
    # If module just passed, Next moves to next module
    if mode == "module_passed":
        # advance module index now
        current_session["module_idx"] += 1
        current_session["mode"] = None
        return update_main_output(start_teaching_module())
    # If idle/none → start teaching module
    if mode is None:
        return update_main_output(start_teaching_module())
    if mode == "teaching":
        return update_main_output(get_explanation())
    if mode == "awaiting_example":
        return update_main_output("✋ Please submit your example using 'Submit Answer' before moving on.")
    if mode == "bloom":
        return update_main_output("🌸 Bloom assessment in progress — answer the Bloom question or submit to retry.")
    if mode == "assessment_summary":
        # show curriculum now and transition to normal teaching flow
        current_session["mode"] = None
        return update_main_output(show_curriculum(current_session["curriculum"]) + "\n\nType 'next' to start Module 1.")
    return update_main_output("⚠️ Invalid state.")

def start_teaching_module():
    if not current_session.get("curriculum"):
        return "⚠️ No curriculum loaded. Please start a session first."
    
    try:
        cur = current_session["curriculum"]
        ch_i = current_session["chapter_idx"]
        m_i = current_session["module_idx"]

        if ch_i >= len(cur.chapters):
            current_session["mode"] = "done"
            return "🎉 You have completed the entire curriculum!"
        
        chapter = cur.chapters[ch_i]
        # if all modules finished -> start Bloom for the chapter
        if m_i >= len(chapter.modules):
            return start_bloom_assessment()

        module = chapter.modules[m_i]
        current_session["mode"] = "teaching"
        return f"📚 Chapter {ch_i+1}: {chapter.name}\nModule {ch_i+1}.{m_i+1}: {module.name}\n\nObjective: {getattr(module,'learning_objective','')}\n\nClick 'Next' to get the explanation."
    except Exception as e:
        return f"❌ Error starting teaching module: {str(e)}"

def get_explanation():
    if not current_session.get("curriculum") or not current_session.get("user"):
        return "⚠️ Session not properly initialized. Please start a new session."
    
    try:
        cur = current_session["curriculum"]
        ch_i = current_session["chapter_idx"]
        m_i = current_session["module_idx"]
        chapter = cur.chapters[ch_i]
        # Guard
        if m_i >= len(chapter.modules):
            return start_bloom_assessment()

        module = chapter.modules[m_i]
        teacher = TeacherAgent(current_session["user"].actual_level or current_session["user"].claimed_level)
        explanation = teacher.teach_module(module)  # returns string
        current_session["mode"] = "awaiting_example"
        return f"📖 Explanation for {module.name}\n\n{explanation}\n\n✍️ Now submit your example in the box and click 'Submit Answer'."
    except Exception as e:
        return f"❌ Error getting explanation: {str(e)}"

# ------------------------------
# Submit handler (single entry point wired to Submit button)
# ------------------------------
def submit_answer(answer):
    mode = current_session.get("mode")
    if mode == "assessment":
        response = handle_assessment(answer)
    elif mode == "awaiting_example":
        response = handle_example_submission(answer)
    elif mode == "bloom":
        response = handle_bloom(answer)
    else:
        response = "⚠️ Nothing to submit right now. Click 'Next' to proceed."
    current_session["last_output"] = response
    # Return output + clear input
    return response, ""

def start_challenge_discussion():
    """
    Opens the challenge discussion chat interface as a modal pop-up.
    Shows the original feedback and prompts user to start discussion.
    """
    mode = current_session.get("mode")
    if mode != "assessment":
        return gr.update(visible=False), [], "⚠️ Challenge is only available during assessment.", get_output_update()
    
    last_fb = current_session.get("last_feedback")
    if not last_fb:
        return gr.update(visible=False), [], "⚠️ No feedback available to challenge. Please submit an answer first.", get_output_update()
    
    # Initialize challenge discussion
    current_session["challenge_mode"] = True
    current_session["challenge_exchanges"] = 0
    current_session["challenge_chat_history"] = []
    
    # Show initial context
    initial_greeting = (
        f"**Grader's Original Feedback:**\n"
        f"Evaluation: {'✅ Correct' if last_fb['is_correct'] else '❌ Incorrect'}\n"
        f"Reasoning: {last_fb['reasoning']}\n\n"
        f"**Question:** {last_fb['question']}\n"
        f"**Your Answer:** {last_fb['answer']}\n\n"
        f"💬 You can now present your arguments. You have up to 3 exchanges with the grader."
    )
    
    # Use messages format: list of dicts with 'role' and 'content'
    chat_history = [{"role": "assistant", "content": initial_greeting}]
    status_msg = "💬 Challenge discussion opened. Present your first argument below (3 exchanges remaining)."
    
    return (
    gr.update(visible=True),
    chat_history,
    status_msg,
    get_output_update()
)

def handle_challenge_message(message, chat_history):
    """
    Handles a message in the challenge discussion.
    Limits to 3 total exchanges (student messages).
    Returns: chat_history, status_msg, msg_enabled, btn_enabled, output_update
    """
    if not current_session.get("challenge_mode"):
        return chat_history, "⚠️ Challenge discussion is not active.", False, False, get_output_update()
    
    if not message or not message.strip():
        return chat_history, "", True, True, gr.update()
    
    if current_session["challenge_exchanges"] >= 3:
        return chat_history, "⚠️ Maximum exchanges (3) reached. Discussion closed. Click 'Close Discussion' to continue.", False, False, get_output_update()
    
    last_fb = current_session.get("last_feedback")
    if not last_fb:
        return chat_history, "⚠️ No feedback available.", False, False, get_output_update()
    
    assessor = LevelAssessmentAgent()
    try:
        grader_response = assessor.challenge_discussion(
            current_session["user"].topic,
            current_session["user"].claimed_level,
            last_fb["question"],
            last_fb["answer"],
            last_fb["evaluation"],
            last_fb["reasoning"],
            current_session["challenge_chat_history"],
            message
        )
        
        current_session["challenge_chat_history"].append([message, grader_response])
        current_session["challenge_exchanges"] += 1
        
        # Ensure chat_history is in messages format (list of dicts)
        chat_history = safe_chat(chat_history)
        
        # Add the new messages in messages format
        chat_history.append({"role": "user", "content": message})
        chat_history.append({"role": "assistant", "content": grader_response})

        
        remaining = 3 - current_session["challenge_exchanges"]
        if remaining > 0:
            status_msg = f"💬 {remaining} exchange(s) remaining. You can continue the discussion."
            output_update = gr.update()
            msg_enabled = True
            btn_enabled = True
        else:
            # Finalize evaluation after 3 exchanges
            status_msg, output_update = finalize_challenge_discussion(grader_response)
            msg_enabled = False
            btn_enabled = False
        
        return chat_history, status_msg, msg_enabled, btn_enabled, output_update
    except Exception as e:
        return chat_history, f"❌ Error: {str(e)}", True, True, gr.update()

def finalize_challenge_discussion(final_response):
    """
    Finalizes the challenge discussion and updates evaluation if needed.
    Extracts final evaluation from the last grader response.
    Returns: (status_msg, main_output_update)
    """
    last_fb = current_session.get("last_feedback")
    if not last_fb:
        return "⚠️ Could not finalize challenge.", get_output_update()
    
    # Try to extract evaluation from the final response
    # Use the challenge_feedback method to get a structured final evaluation
    assessor = LevelAssessmentAgent()
    try:
        # Get the full conversation context
        conversation_text = "\n".join([
            f"Student: {msg[0]}\nGrader: {msg[1]}"
            for msg in current_session["challenge_chat_history"]
        ])
        
        # Final re-evaluation request
        final_prompt = (
            f"Based on our discussion:\n{conversation_text}\n\n"
            "Please provide your FINAL evaluation as JSON with: "
            '{"evaluation": "correct" or "incorrect", "reasoning": "explanation", "original_was_fair": true/false}'
        )
        
        raw = assessor.challenge_feedback(
            current_session["user"].topic,
            current_session["user"].claimed_level,
            last_fb["question"],
            last_fb["answer"],
            last_fb["evaluation"],
            last_fb["reasoning"]
        )
        
        import json as _json
        parsed = _json.loads(raw)
        new_evaluation = str(parsed.get("evaluation", "incorrect")).lower()
        new_reasoning = parsed.get("reasoning", "")
        original_was_fair = parsed.get("original_was_fair", True)
        
        new_is_correct = new_evaluation == "correct"
        old_is_correct = last_fb["is_correct"]
        
        # Build the main output update
        qs = current_session["questions"]
        idx = current_session["q_idx"]
        current_question_text = ""
        if idx < len(qs):
            current_question_text = f"\n\n📝 Current Question {idx+1} of {len(qs)}:\n{qs[idx]}"
        
        # Update if evaluation changed
        if new_is_correct != old_is_correct:
            if new_is_correct and not old_is_correct:
                current_session["correct_count"] += 1
            elif not new_is_correct and old_is_correct:
                current_session["correct_count"] = max(0, current_session["correct_count"] - 1)
            
            if current_session["assessment_feedback"]:
                current_session["assessment_feedback"][-1]["correct"] = new_is_correct
                current_session["assessment_feedback"][-1]["reason"] = new_reasoning
            
            current_session["last_feedback"]["evaluation"] = new_evaluation
            current_session["last_feedback"]["reasoning"] = new_reasoning
            current_session["last_feedback"]["is_correct"] = new_is_correct
            
            # Build updated main output
            main_output = (
                f"🔄 **EVALUATION UPDATED AFTER CHALLENGE**\n\n"
                f"**Question:** {last_fb['question']}\n"
                f"**Your Answer:** {last_fb['answer']}\n\n"
                f"**Original Evaluation:** {'✅ Correct' if old_is_correct else '❌ Incorrect'}\n"
                f"**Updated Evaluation:** {'✅ Correct' if new_is_correct else '❌ Incorrect'}\n\n"
                f"**Updated Reasoning:** {new_reasoning}\n"
                f"{current_question_text}"
            )
            
            status_msg = "🔄 **FINAL EVALUATION UPDATED**\n\n" + \
                        f"**New Evaluation:** {'✅ Correct' if new_is_correct else '❌ Incorrect'}\n" + \
                        f"**Final Reasoning:** {new_reasoning}\n\n" + \
                        "✅ Challenge discussion completed. Evaluation has been updated on the main screen."
        else:
            # Build main output showing final evaluation
            main_output = (
                f"📋 **FINAL EVALUATION AFTER CHALLENGE**\n\n"
                f"**Question:** {last_fb['question']}\n"
                f"**Your Answer:** {last_fb['answer']}\n\n"
                f"**Evaluation:** {'✅ Correct' if new_is_correct else '❌ Incorrect'} (unchanged)\n\n"
                f"**Final Reasoning:** {new_reasoning}\n"
                f"{current_question_text}"
            )
            
            status_msg = "📋 **FINAL EVALUATION**\n\n" + \
                        f"**Evaluation:** {'✅ Correct' if new_is_correct else '❌ Incorrect'} (unchanged)\n" + \
                        f"**Final Reasoning:** {new_reasoning}\n\n" + \
                        "✅ Challenge discussion completed. You may continue with the assessment."
        
        # Close challenge mode
        current_session["challenge_mode"] = False
        current_session["last_output"] = main_output
        
        return status_msg, gr.update(value=main_output)
    except Exception as e:
        current_session["challenge_mode"] = False
        error_msg = f"⚠️ Error finalizing challenge: {str(e)}"
        return error_msg, get_output_update()

def handle_example_submission(example_text):
    """
    Uses TeacherAgent.evaluate_example(module, example) to decide correctness.
    If correct -> set mode to 'module_passed' and require user to press Next to move on.
    If incorrect -> remain in 'awaiting_example' and show feedback.
    """
    if not current_session.get("curriculum"):
        return "⚠️ No curriculum loaded. Please start a session first."
    
    cur = current_session["curriculum"]
    ch_i = current_session["chapter_idx"]
    m_i = current_session["module_idx"]
    
    if ch_i >= len(cur.chapters):
        return "⚠️ No active chapter available."
    
    chapter = cur.chapters[ch_i]
    if m_i >= len(chapter.modules):
        # Shouldn't happen, but guard
        return "⚠️ No active module to evaluate."

    module = chapter.modules[m_i]
    teacher = TeacherAgent(current_session["user"].actual_level or current_session["user"].claimed_level)

    # quick client-side guardrails for empty/brief examples
    if not example_text or not str(example_text).strip():
        return "❌ Please provide an example to demonstrate your understanding."
    if len(str(example_text).strip().split()) < 10:
        return "❌ Your example is too brief. Provide 2-3 sentences with specific details."

    try:
        eval_result = teacher.evaluate_example(module, example_text)
        is_correct = bool(eval_result.get("is_correct"))
        feedback = eval_result.get("feedback", "No feedback provided.")
        confidence = eval_result.get("confidence", None)
    except Exception as e:
        return f"❌ Error evaluating example: {str(e)}. Please try again."

    if is_correct:
        # mark module as passed (do not auto-increment module_idx — require Next)
        current_session["mode"] = "module_passed"
        return f"✅ Example accepted. Feedback: {feedback}\n\n➡️ Click 'Next' to continue to the next module."
    else:
        # remain in awaiting_example — must retry
        current_session["mode"] = "awaiting_example"
        return f"❌ Example not sufficient. Feedback: {feedback}\n\nPlease try another example for the same module."

# ------------------------------
# Bloom assessment (per chapter)
# ------------------------------
def start_bloom_assessment():
    if not current_session.get("curriculum"):
        return "⚠️ No curriculum loaded. Please start a session first."
    
    try:
        current_session["mode"] = "bloom"
        current_session["bloom_level"] = BLOOM_ORDER[0]
        chapter = current_session["curriculum"].chapters[current_session["chapter_idx"]]
        return ask_bloom_question(current_session["bloom_level"], chapter)
    except Exception as e:
        return f"❌ Error starting Bloom assessment: {str(e)}"

def ask_bloom_question(level, chapter):
    try:
        agent = BloomsAssessmentAgent()
        # agent.generate_bloom_question expects (chapter, bloom_level)
        q = agent.generate_bloom_question(chapter, level)
        current_session["questions"] = [q]
        return f"🌸 Bloom's Assessment ({level.title()}) for Chapter: {chapter.name}\n\n{q}\n\n✍️ Answer below and click Submit."
    except Exception as e:
        return f"❌ Error generating Bloom question: {str(e)}"

def handle_bloom(answer):
    if not current_session.get("curriculum"):
        return "⚠️ No curriculum loaded. Please start a session first."
    
    if not answer or not str(answer).strip():
        return "❌ Please provide an answer for the Bloom assessment question."
    
    try:
        level = current_session["bloom_level"]
        if not current_session["questions"]:
            return "⚠️ No question available. Please start a new session."
        question = current_session["questions"][0]
        chapter = current_session["curriculum"].chapters[current_session["chapter_idx"]]
        agent = BloomsAssessmentAgent()

        # agent.evaluate_bloom_answer returns a JSON string per your agent implementation
        raw_eval = agent.evaluate_bloom_answer(question, answer, level, chapter)

        # parse evaluation JSON
        try:
            eval_obj = json.loads(raw_eval)
            score = float(eval_obj.get("score", 0))
            feedback = str(eval_obj.get("feedback", "No feedback"))
        except (json.JSONDecodeError, ValueError) as e:
            # fallback: if text contains 'correct' treat as pass
            feedback = f"Could not parse evaluation response. Raw: {raw_eval[:100]}..."
            score = 10 if "correct" in raw_eval.lower() else 0
    except KeyError as e:
        return f"❌ Session error: Missing required data ({str(e)}). Please restart the session."
    except Exception as e:
        return f"❌ Error evaluating Bloom answer: {str(e)}. Please try again."

    # use threshold (e.g., >=6/10)
    if score >= 6:
        # advance bloom level
        next_idx = BLOOM_ORDER.index(level) + 1
        if next_idx < len(BLOOM_ORDER):
            current_session["bloom_level"] = BLOOM_ORDER[next_idx]
            # generate new question for next level
            return f"✅ {feedback}\n\n➡️ Moving to {BLOOM_ORDER[next_idx].title()}.\n\n" + ask_bloom_question(current_session["bloom_level"], chapter)
        else:
            # finished Bloom for chapter -> next chapter
            current_session["chapter_idx"] += 1
            current_session["module_idx"] = 0
            current_session["mode"] = None
            return f"🎉 {feedback}\n\n✅ Bloom’s assessment completed for chapter '{chapter.name}'.\nType 'next' to continue."
    else:
        # ask a new question for same level
        return f"❌ {feedback}\n\n🔁 Try another question at the same level.\n\n" + ask_bloom_question(level, chapter)

# ------------------------------
# Gradio UI wiring (single unified output)
# ------------------------------
# Custom CSS for modal pop-up
modal_css = """
.modal-overlay:not([style*="display: none"]) {
    position: fixed !important;
    top: 0 !important;
    left: 0 !important;
    width: 100% !important;
    height: 100% !important;
    background-color: rgba(0, 0, 0, 0.5) !important;
    z-index: 1000 !important;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
    padding: 20px !important;
}

/* When Gradio hides the element, ensure it doesn't block interactions */
.modal-overlay[style*="display: none"],
.modal-overlay[style*="display:none"] {
    display: none !important;
    visibility: hidden !important;
    pointer-events: none !important;
    z-index: -1 !important;
    opacity: 0 !important;
}

.modal-content {
    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%) !important;
    border-radius: 10px !important;
    padding: 20px !important;
    max-width: 900px !important;
    width: 100% !important;
    max-height: 90vh !important;
    overflow-y: auto !important;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.3) !important;
}

/* Chat window styling */
.modal-content .gradio-chatbot {
    background-color: #ffffff !important;
    border-radius: 8px !important;
    padding: 15px !important;
    border: 2px solid #e0e0e0 !important;
}

.modal-content .gradio-chatbot .message {
    background-color: #f8f9fa !important;
    border-radius: 8px !important;
    padding: 10px !important;
    margin: 5px 0 !important;
}

.modal-content .gradio-chatbot .user-message {
    background-color: #e3f2fd !important;
    border-left: 4px solid #2196f3 !important;
}

.modal-content .gradio-chatbot .bot-message {
    background-color: #f1f8e9 !important;
    border-left: 4px solid #8bc34a !important;
}

.modal-header {
    margin: 0 !important;
    padding: 0 !important;
    flex-grow: 1 !important;
}

.modal-close-btn {
    min-width: 40px !important;
    height: 40px !important;
    border-radius: 50% !important;
    font-size: 20px !important;
    font-weight: bold !important;
}

.modal-note {
    font-size: 12px !important;
    color: #666 !important;
    margin-top: 10px !important;
}
"""

with gr.Blocks(title="AI Teacher Bot") as demo:
    gr.Markdown("""
# 🧠 AI Teacher Bot

In this interactive learning experience, you will be prompted to select your learning level—**Beginner**, **Intermediate**, or **Advanced**. 

- If you choose **Intermediate** or **Advanced**, you will be presented with an assessment designed to test your understanding of the material. The questions in these assessments are tailored to the selected level, ensuring they are **challenging and reflective of the knowledge expected at that stage**. For instance, **Intermediate-level questions** will require more in-depth explanations, not just basic one-liner answers. Similarly, **Advanced-level assessments** will be comprehensive and demand a higher level of critical thinking and subject mastery.
""")


    with gr.Row():
        topic = gr.Textbox(label="📝 Topic", placeholder="e.g., Python Programming", scale=2)
        level = gr.Dropdown(choices=LEVELS, value="novice", label="🎓 Your Level", scale=1)
        start_btn = gr.Button("🚀 Start Session")

    output = gr.Textbox(label="📚 Session Output", lines=25, interactive=False, autoscroll=True)
    answer_box = gr.Textbox(label="✍️ Your Answer / Example", placeholder="Type your answer...", lines=5, max_lines=10)
    with gr.Row():
        submit_btn = gr.Button("Submit Answer")
        next_btn = gr.Button("Next")
        challenge_btn = gr.Button("Challenge Assessment", visible=False)
    
    # Challenge discussion chat interface - Modal Pop-up
    with gr.Column(visible=False, elem_classes="modal-overlay") as challenge_modal_overlay:
        with gr.Column(elem_classes="modal-content"):
            with gr.Row():
                gr.Markdown("### 💬 Challenge Discussion with Grader", elem_classes="modal-header")
                challenge_close_btn = gr.Button("✕", elem_classes="modal-close-btn", scale=0)
            
            with gr.Row():
                with gr.Column(scale=3):
                    challenge_chat = gr.Chatbot(
                        label="",
                        height=400,
                        show_label=False,
                        container=True
                    )
                    challenge_status = gr.Textbox(
                        label="Status",
                        interactive=False,
                        lines=2,
                        container=True
                    )
                with gr.Column(scale=1):
                    challenge_msg_box = gr.Textbox(
                        label="Your Argument/Question",
                        placeholder="Explain why you think your answer is correct...",
                        lines=5,
                        container=True
                    )
                    challenge_send_btn = gr.Button("Send", variant="primary")
                    gr.Markdown("**Note:** You have up to 3 exchanges with the grader.", elem_classes="modal-note")

    # Function to update challenge button visibility
    def update_challenge_visibility():
        mode = current_session.get("mode")
        has_feedback = current_session.get("last_feedback") is not None
        # Show button during assessment mode after first answer is submitted
        should_show = (mode == "assessment" and has_feedback)
        return gr.update(visible=should_show)
    
    # handlers
    start_btn.click(
        fn=start_learning_session, 
        inputs=[topic, level], 
        outputs=[output]
    ).then(
        fn=update_challenge_visibility,
        inputs=None,
        outputs=[challenge_btn]
    )
    
    submit_btn.click(
        fn=submit_answer, 
        inputs=[answer_box], 
        outputs=[output, answer_box]
    ).then(
        fn=update_challenge_visibility,
        inputs=None,
        outputs=[challenge_btn]
    )
    
    next_btn.click(
        fn=next_step, 
        inputs=[answer_box], 
        outputs=[output]
    ).then(
        fn=update_challenge_visibility,
        inputs=None,
        outputs=[challenge_btn]
    )
    challenge_btn.click(
        fn=start_challenge_discussion,
        inputs=None,
        outputs=[challenge_modal_overlay, challenge_chat, challenge_status, output]
    ).then(
        fn=lambda: (gr.update(value="", interactive=True), gr.update(interactive=True)),
        inputs=None,
        outputs=[challenge_msg_box, challenge_send_btn]
    )
    
    def send_and_clear(message, chat_history):
        """Send message and clear input box (if needed)"""
        history, status, msg_enabled, btn_enabled, output_update = handle_challenge_message(message, chat_history)
        if message and message.strip():
            msg_update = gr.update(value="", interactive=msg_enabled)
        else:
            msg_update = gr.update(value=message, interactive=msg_enabled)
        btn_update = gr.update(interactive=btn_enabled)
        return history, msg_update, status, btn_update, output_update
    
    challenge_send_btn.click(
        fn=send_and_clear,
        inputs=[challenge_msg_box, challenge_chat],
        outputs=[challenge_chat, challenge_msg_box, challenge_status, challenge_send_btn, output]
    )
    
    challenge_msg_box.submit(
        fn=send_and_clear,
        inputs=[challenge_msg_box, challenge_chat],
        outputs=[challenge_chat, challenge_msg_box, challenge_status, challenge_send_btn, output]
    )
    
    def close_challenge_discussion():
        """
        Properly closes the challenge modal and restores UI control
        """
        # Reset backend state
        current_session["challenge_mode"] = False
        current_session["challenge_chat_history"] = []
        current_session["challenge_exchanges"] = 0

        # IMPORTANT: Use visible=False to hide the overlay completely
        # This should remove it from the DOM or set display:none which CSS will respect
        return (
            gr.update(visible=False),   # Hide overlay - this should remove blocking
            empty_chat(),               # Reset chat
            "✅ Challenge closed.",     # Status message
            gr.update(value=current_session.get("last_output", ""))  # Keep main output
        )

    
    challenge_close_btn.click(
        fn=close_challenge_discussion,
        inputs=None,
        outputs=[
            challenge_modal_overlay,  # visible=False
            challenge_chat,           # reset chat
            challenge_status,         # status text
            output                    # main output (reevaluated decision)
        ]
    ).then(
        fn=update_challenge_visibility,
        inputs=None,
        outputs=[challenge_btn]
    )


if __name__ == "__main__":
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860,
        theme=gr.themes.Soft(),
        css=modal_css
    )