# app_AD2class.py
# Main application file. Builds a tabbed Gradio UI for multiple assessments.

import gradio as gr

# Import the separated modules
import MMSE
import CDRM  # <-- Import the new placeholder module
import utils

# --- MMSE UI and Logic ---
def build_mmse_tab(mmse_score_state):
    """Builds the UI components for the MMSE tab."""
    with gr.Blocks() as mmse_tab:
        gr.Markdown("## Mini Mental State Exam")
        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']} - Q{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=(initial_q['main_num'] == 11))

        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=15, interactive=False)
        total = gr.Textbox(label="Total Score", interactive=False)

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

        # --- Event Handler Functions specific to this tab ---
        def process_evaluation_MMSE(answers, image_upload):
            score_lines_out, total_out = MMSE.evaluate_MMSE(answers, image_upload)
            return score_lines_out, total_out, gr.update(visible=False), gr.update(visible=True), total_out
        
        # Wire up 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(lambda *args: utils.save_and_navigate("next", *args, module=MMSE), 
                       inputs=[question_index, answer_text, answers], outputs=outputs_list)
        prev_btn.click(lambda *args: utils.save_and_navigate("prev", *args, module=MMSE), 
                       inputs=[question_index, answer_text, answers], outputs=outputs_list)
        jump_nav.change(lambda *args: utils.jump_to_question(*args, module=MMSE), 
                        inputs=[jump_nav, question_index, answer_text, answers], outputs=outputs_list)

        submit_btn.click(
            utils.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, mmse_score_state]
        )
        
        start_over_btn.click(lambda: utils.reset_app(MMSE), 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, mmse_score_state
        ])

    return mmse_tab


# --- CDRM UI and Logic (Placeholder) ---
def build_cdrm_tab(cdrm_score_state):
    """Builds the fully functional UI components for the CDR Memory tab."""
    with gr.Blocks() as cdrm_tab:
        gr.Markdown("## CDR Memory Assessment")
        # State and UI components, mirroring the MMSE tab
        question_index = gr.State(0)
        answers = gr.State([""] * CDRM.TOTAL_QUESTIONS)

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

        initial_q = CDRM.STRUCTURED_QUESTIONS[0]
        progress_text = gr.Markdown(f"## {initial_q['main_cat']} - Q{initial_q['main_num']}{initial_q['sub_letter']} ({1} of {CDRM.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")
            # This is now dynamic; it will only show if a drawing question is added to CDRM.py
            image_upload = gr.Image(type="filepath", label="Upload Drawing", visible=("draw a copy" in initial_q["question"]))

        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=15, interactive=False)
        total = gr.Textbox(label="Total Score", interactive=False)

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

        # --- Event Handler Functions specific to this tab ---
        def process_evaluation_CDRM(answers):
            # This now passes the answers to the evaluation function
            score_lines_out, total_out = CDRM.evaluate_CDRM(answers)
            return score_lines_out, total_out, gr.update(visible=False), gr.update(visible=True), total_out

        # Wire up event listeners, referencing the CDRM module
        audio_input.change(utils.transcribe, inputs=audio_input, outputs=answer_text)
        question_button.click(CDRM.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]

        # Use the generic utils functions, passing 'module=CDRM'
        next_btn.click(lambda *args: utils.save_and_navigate("next", *args, module=CDRM),
                       inputs=[question_index, answer_text, answers], outputs=outputs_list)
        prev_btn.click(lambda *args: utils.save_and_navigate("prev", *args, module=CDRM),
                       inputs=[question_index, answer_text, answers], outputs=outputs_list)
        jump_nav.change(lambda *args: utils.jump_to_question(*args, module=CDRM),
                        inputs=[jump_nav, question_index, answer_text, answers], outputs=outputs_list)

        submit_btn.click(
            utils.save_final_answer, inputs=[question_index, answer_text, answers], outputs=answers
        ).then(
            fn=process_evaluation_CDRM, inputs=[answers], outputs=[score_lines, total, submit_btn, start_over_btn, cdrm_score_state]
        )

        # The reset function also takes the module as an argument
        start_over_btn.click(lambda: utils.reset_app(CDRM), 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, cdrm_score_state
        ])

    return cdrm_tab


# --- Joint Outcome UI ---
def build_joint_outcome_tab(mmse_score_state, cdrm_score_state):
    """Builds the UI for the Joint Outcome analysis tab."""
    with gr.Blocks() as joint_tab:
        gr.Markdown("## Joint Outcome Analysis")
        gr.Markdown("This tab will use the results from the MMSE and CDR Memory tabs to provide a combined analysis.")
        
        with gr.Row():
            mmse_result = gr.Textbox(label="MMSE Final Score", interactive=False)
            cdrm_result = gr.Textbox(label="CDR Memory Final Score", interactive=False)

        analyze_btn = gr.Button("Analyze Joint Outcome")
        outcome_display = gr.Textbox(label="Combined Analysis Result", interactive=False)

        def analyze_outcomes(mmse_score, cdrm_score):
            # Placeholder logic for joint analysis
            if not mmse_score and not cdrm_score:
                return "Please complete both assessments first."
            
            # Example:
            # mmse_val = int(mmse_score.split('/')[0].strip())
            # cdrm_val = int(cdrm_score.split('/')[0].strip())
            # some_logic...
            
            return f"Analysis based on MMSE score: '{mmse_score}' and CDRM score: '{cdrm_score}'.\n\n(Full analysis logic to be implemented)"

        analyze_btn.click(analyze_outcomes, inputs=[mmse_score_state, cdrm_score_state], outputs=outcome_display)
        
        # Link state changes to the textboxes for display
        mmse_score_state.change(lambda x: x, inputs=mmse_score_state, outputs=mmse_result)
        cdrm_score_state.change(lambda x: x, inputs=cdrm_score_state, outputs=cdrm_result)

    return joint_tab

# --- Main Application Build ---
def build_ui():
    """Builds the main Gradio interface with tabs."""
    with gr.Blocks(theme=gr.themes.Soft(), title="Cognitive Assessment") as demo:
        gr.Markdown("# Cognitive Assessment")
        
        # State objects to share results between tabs
        mmse_score_state = gr.State("")
        cdrm_score_state = gr.State("")
        
        with gr.Tabs():
            with gr.TabItem("Joint Outcome"):
                build_joint_outcome_tab(mmse_score_state, cdrm_score_state)
            with gr.TabItem("MMSE"):
                build_mmse_tab(mmse_score_state)
            with gr.TabItem("CDR Memory"):
                build_cdrm_tab(cdrm_score_state)

    return demo

if __name__ == "__main__":
    # Perform one-time setup for all modules
    MMSE.pregenerate_audio()
    CDRM.pregenerate_audio()
    
    app = build_ui()
    app.launch()