# app_ADclass.py
# Main application file. Focuses on building the Gradio UI and handling
# user interactions. It imports assessment logic from MMSE.py and utilities from utils.py.

import gradio as gr

# Import the separated modules
import MMSE
import utils

# --- UI Interaction Handlers ---
def update_view(new_index, all_answers):
    """Updates the UI elements when navigating to a new question."""
    q_data = MMSE.STRUCTURED_QUESTIONS[new_index]
    progress = f"## {q_data['main_cat']} - Question {q_data['main_num']}{q_data['sub_letter']} ({new_index + 1} of {MMSE.TOTAL_QUESTIONS})"
    drawing_q_visible = (new_index == MMSE.DRAWING_Q_INDEX)

    return (
        f"Say 🔊 {q_data['question']}",
        all_answers[new_index],
        new_index,
        progress,
        q_data["instruction"],
        MMSE.QUESTION_CHOICES[new_index],
        gr.update(visible=drawing_q_visible),
        None  # Clear the audio_input component
    )

def save_and_navigate(direction, current_index, current_answer, all_answers):
    """Saves the current answer and moves to the next/previous question."""
    all_answers[current_index] = current_answer
    if direction == "next":
        new_index = min(current_index + 1, MMSE.TOTAL_QUESTIONS - 1)
    else:  # prev
        new_index = max(current_index - 1, 0)
    return update_view(new_index, all_answers) + (all_answers,)

def jump_to_question(selected_choice, current_index, current_answer, all_answers):
    """Saves the current answer and jumps to the selected question."""
    if not selected_choice:
        return update_view(current_index, all_answers) + (all_answers,)
    all_answers[current_index] = current_answer
    new_index = MMSE.QUESTION_CHOICES.index(selected_choice)
    return update_view(new_index, all_answers) + (all_answers,)

def save_final_answer(current_index, current_answer, all_answers):
    """A dedicated function to save the last answer before submitting."""
    all_answers[current_index] = current_answer
    return all_answers

def process_evaluation_MMSE(answers, image_upload):
    """Calls the evaluation function from the MMSE module and prepares UI updates."""
    score_lines, total = MMSE.evaluate_MMSE(answers, image_upload)
    # This function now returns the data AND the UI updates
    return score_lines, total, gr.update(visible=False), gr.update(visible=True)


def reset_app():
    """Resets the entire application state to its initial view."""
    initial_q = MMSE.STRUCTURED_QUESTIONS[0]
    return (
        0,  # question_index
        [""] * MMSE.TOTAL_QUESTIONS,  # answers
        "",  # score_lines
        "",  # total
        f"Say 🔊 {initial_q['question']}",  # question_button
        f"## {initial_q['main_cat']} - Question {initial_q['main_num']}{initial_q['sub_letter']} (1 of {MMSE.TOTAL_QUESTIONS})",  # progress_text
        initial_q["instruction"],  # instruction_display
        "",  # answer_text
        MMSE.QUESTION_CHOICES[0],  # jump_nav
        None,  # audio_input
        None,  # image_upload
        gr.update(visible=False),  # start_over_btn
        gr.update(visible=True),   # submit_btn
        None   # tts_audio
    )

def build_ui():
    """Builds and returns the Gradio interface."""
    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown("# Mini Mental State Exam (MMSE) with Multimodal Inputs")
        question_index = gr.State(0)
        answers = gr.State([""] * MMSE.TOTAL_QUESTIONS)

        jump_nav = gr.Dropdown(choices=MMSE.QUESTION_CHOICES, value=MMSE.QUESTION_CHOICES[0], label="Jump to Question")

        initial_q = MMSE.STRUCTURED_QUESTIONS[0]
        progress_text = gr.Markdown(f"## {initial_q['main_cat']} - Question {initial_q['main_num']}{initial_q['sub_letter']} (1 of {MMSE.TOTAL_QUESTIONS})")
        instruction_display = gr.Markdown(initial_q["instruction"])
        question_button = gr.Button(f"Say 🔊 {initial_q['question']}", variant="secondary")

        with gr.Group():
            answer_text = gr.Textbox(label="Your Answer (type or record below)")
            audio_input = gr.Audio(sources=["microphone"], label="🎤 Record Your Answer Here")
            image_upload = gr.Image(type="filepath", label="Upload Drawing for Drawing Question", visible=False)

        with gr.Row():
            prev_btn = gr.Button("⬅️ Previous Question")
            next_btn = gr.Button("Next Question ➡️")

        tts_audio = gr.Audio(autoplay=True, visible=False)
        score_lines = gr.Textbox(label="Score by Question", lines=MMSE.TOTAL_QUESTIONS)
        total = gr.Textbox(label="Total Score")

        submit_btn = gr.Button("✅ Submit All Answers")
        start_over_btn = gr.Button("🔄 Start Over", visible=False)

        # --- Event Listeners ---
        audio_input.change(utils.transcribe, inputs=audio_input, outputs=answer_text)
        question_button.click(MMSE.speak_question, inputs=question_index, outputs=tts_audio)

        outputs_list = [question_button, answer_text, question_index, progress_text,
                        instruction_display, jump_nav, image_upload, audio_input, answers]

        next_btn.click(save_and_navigate, inputs=[gr.Textbox("next", visible=False), question_index, answer_text, answers], outputs=outputs_list)
        prev_btn.click(save_and_navigate, inputs=[gr.Textbox("prev", visible=False), question_index, answer_text, answers], outputs=outputs_list)
        jump_nav.change(jump_to_question, inputs=[jump_nav, question_index, answer_text, answers], outputs=outputs_list)

        submit_btn.click(
            save_final_answer, inputs=[question_index, answer_text, answers], outputs=answers
        ).then(
            fn=process_evaluation_MMSE, inputs=[answers, image_upload], outputs=[score_lines, total, submit_btn, start_over_btn]
        )
        
        reset_outputs = [question_index, answers, score_lines, total, question_button, progress_text,
                         instruction_display, answer_text, jump_nav, audio_input, image_upload,
                         start_over_btn, submit_btn, tts_audio]
        start_over_btn.click(reset_app, inputs=None, outputs=reset_outputs)

    return demo

if __name__ == "__main__":
    # Perform one-time setup before launching the UI
    MMSE.pregenerate_audio()
    
    # Build and launch the UI
    app = build_ui()
    app.launch()