import gradio as gr
from dotenv import load_dotenv
from .objective_handlers import (
    process_files,
    process_files_and_generate_questions,
    process_user_objectives,
    process_user_objectives_and_generate_questions,
)
from .question_handlers import generate_questions
from .edit_handlers import load_quiz_for_editing, load_file_for_editing, accept_and_next, go_previous, save_and_download
from .formatting import format_quiz_for_ui
from .run_manager import get_run_manager
from models import MODELS

# Load environment variables
load_dotenv()

# Set to False to disable saving output files (folders, logs, JSON, markdown).
# Tab 3 download of edited questions still works.
SAVE_OUTPUTS = True

def create_ui():
    """Create the Gradio UI."""
    get_run_manager().save_outputs = SAVE_OUTPUTS
    with gr.Blocks(title="Quiz Generator V3") as app:
        gr.Markdown("# Quiz Generator V3")
        gr.Markdown("Upload course materials and generate learning objectives and quiz questions.")

        with gr.Tab("Generate Learning Objectives"):
            with gr.Row():
                with gr.Column():
                    mode_radio = gr.Radio(
                        choices=["Generate from course materials", "Use my own learning objectives"],
                        value="Generate from course materials",
                        label="Learning Objectives Source"
                    )
                    files_input = gr.File(
                        file_count="multiple",
                        label="Upload Course Materials (.vtt, .srt, .ipynb, .md)",
                        file_types=[".ipynb", ".vtt", ".srt", ".md"],
                        type="filepath"
                    )
                    # Generate-mode controls
                    num_objectives = gr.Slider(minimum=1, maximum=20, value=4, step=1, label="Number of Learning Objectives per Run")
                    # Manual-mode input (hidden by default)
                    user_objectives_input = gr.Textbox(
                        label="Enter Learning Objectives (one per line)",
                        lines=8,
                        placeholder=(
                            "Enter each learning objective on a new line.\n"
                            "Leading numbers or letters are stripped automatically.\n\n"
                            "Example:\n"
                            "1. Identify the key stages of the data engineering lifecycle\n"
                            "2. Articulate a mental framework for building solutions\n"
                            "3. Identify upstream and downstream stakeholders"
                        ),
                        visible=False
                    )
                    with gr.Accordion("Advanced Options", open=False):
                        num_runs = gr.Dropdown(
                            choices=["1", "2", "3", "4", "5"],
                            value="2",
                            label="Number of Generation Runs"
                        )
                        model_dropdown = gr.Dropdown(
                            choices=MODELS,
                            value="gpt-5.2",
                            label="Model"
                        )
                        incorrect_answer_model_dropdown = gr.Dropdown(
                            choices=MODELS,
                            value="gpt-5.2",
                            label="Model for Incorrect Answer Suggestions"
                        )
                        temperature_dropdown = gr.Dropdown(
                            choices=["0.0", "0.1", "0.2", "0.3", "0.4", "0.5", "0.6", "0.7", "0.8", "0.9", "1.0"],
                            value="1.0",
                            label="Temperature (0.0: Deterministic, 1.0: Creative)"
                        )
                    # Buttons — generate_button / process_user_obj_button toggle by mode;
                    # generate_all_button is always visible
                    generate_button = gr.Button("Generate Learning Objectives")
                    process_user_obj_button = gr.Button("Process Learning Objectives", visible=False)
                    generate_all_button = gr.Button("Generate all", variant="primary")

                with gr.Column():
                    status_output = gr.Textbox(label="Status")
                    objectives_output = gr.Textbox(label="Best-in-Group Learning Objectives", lines=10)
                    grouped_output = gr.Textbox(label="All Grouped Learning Objectives", lines=10)
                    raw_ungrouped_output = gr.Textbox(label="Raw Ungrouped Learning Objectives (Debug)", lines=10)

        with gr.Tab("Generate Questions"):
            with gr.Row():
                with gr.Column():
                    objectives_input = gr.Textbox(label="Learning Objectives JSON", lines=10, max_lines=10)
                    num_questions_slider = gr.Slider(minimum=1, maximum=10, value=10, step=1, label="Number of questions")
                    with gr.Accordion("Advanced Options", open=False):
                        model_dropdown_q = gr.Dropdown(
                            choices=MODELS,
                            value="gpt-5.2",
                            label="Model"
                        )
                        temperature_dropdown_q = gr.Dropdown(
                            choices=["0.0", "0.1", "0.2", "0.3", "0.4", "0.5", "0.6", "0.7", "0.8", "0.9", "1.0"],
                            value="1.0",
                            label="Temperature (0.0: Deterministic, 1.0: Creative)"
                        )
                        num_runs_q = gr.Slider(minimum=1, maximum=5, value=2, step=1, label="Number of Question Generation Runs")
                    generate_q_button = gr.Button("Generate Questions")

                with gr.Column():
                    status_q_output = gr.Textbox(label="Status")
                    best_questions_output = gr.Textbox(label="Ranked Best-in-Group Questions", lines=10)
                    all_questions_output = gr.Textbox(label="All Grouped Questions", lines=10)
                    formatted_quiz_output = gr.Textbox(label="Formatted Quiz", lines=15)

        with gr.Tab("Propose/Edit Question"):
            # State for editing flow
            questions_state = gr.State([])
            index_state = gr.State(0)
            edited_state = gr.State([])

            with gr.Row():
                with gr.Column():
                    edit_status = gr.Textbox(label="Status", interactive=False)
                    with gr.Accordion("Load questions from file (.md or .yml)", open=False):
                        file_upload = gr.File(
                            label="Upload a quiz file (.md or .yml)",
                            file_types=[".md", ".yml", ".yaml"],
                            type="filepath"
                        )
                        load_file_button = gr.Button("Load from file")
                    edit_button = gr.Button("Edit questions from Tab 2", variant="primary")
                    question_editor = gr.Textbox(
                        label="Question",
                        lines=15,
                        interactive=True,
                        placeholder="Click 'Edit questions' to load the generated quiz."
                    )
                    with gr.Row():
                        prev_button = gr.Button("Previous")
                        next_button = gr.Button("Accept & Next", variant="primary")
                    download_button = gr.Button("Download edited quiz")
                    download_file = gr.File(label="Download", interactive=False)

        # --- Mode toggle: show/hide controls based on selected source ---
        def _toggle_objective_mode(mode):
            is_generate = mode == "Generate from course materials"
            return (
                gr.update(visible=is_generate),       # num_objectives
                gr.update(visible=is_generate),       # num_runs
                gr.update(visible=not is_generate),   # user_objectives_input
                gr.update(visible=is_generate),       # generate_button
                gr.update(visible=not is_generate),   # process_user_obj_button
            )

        mode_radio.change(
            _toggle_objective_mode,
            inputs=[mode_radio],
            outputs=[num_objectives, num_runs, user_objectives_input,
                     generate_button, process_user_obj_button]
        )

        # --- "Generate all" dispatcher: routes based on current mode ---
        def _generate_all(mode, files, num_obj, num_r, user_obj_text,
                          model, ia_model, temp,
                          model_q, temp_q, num_q, num_runs_q_val):
            if mode == "Generate from course materials":
                return process_files_and_generate_questions(
                    files, num_obj, num_r, model, ia_model, temp,
                    model_q, temp_q, num_q, num_runs_q_val
                )
            else:
                return process_user_objectives_and_generate_questions(
                    files, user_obj_text, model, ia_model, temp,
                    model_q, temp_q, num_q, num_runs_q_val
                )

        # Set up event handlers
        generate_button.click(
            process_files,
            inputs=[files_input, num_objectives, num_runs, model_dropdown, incorrect_answer_model_dropdown, temperature_dropdown],
            outputs=[status_output, objectives_output, grouped_output, raw_ungrouped_output]
        )

        process_user_obj_button.click(
            process_user_objectives,
            inputs=[files_input, user_objectives_input, model_dropdown, incorrect_answer_model_dropdown, temperature_dropdown],
            outputs=[status_output, objectives_output, grouped_output, raw_ungrouped_output]
        )

        generate_all_button.click(
            _generate_all,
            inputs=[
                mode_radio,
                files_input, num_objectives, num_runs, user_objectives_input,
                model_dropdown, incorrect_answer_model_dropdown, temperature_dropdown,
                model_dropdown_q, temperature_dropdown_q, num_questions_slider, num_runs_q
            ],
            outputs=[
                status_output, objectives_output, grouped_output, raw_ungrouped_output,
                status_q_output, best_questions_output, all_questions_output, formatted_quiz_output
            ]
        )

        objectives_output.change(
            lambda x: x,
            inputs=[objectives_output],
            outputs=[objectives_input]
        )

        generate_q_button.click(
            generate_questions,
            inputs=[objectives_input, model_dropdown_q, temperature_dropdown_q, num_questions_slider, num_runs_q],
            outputs=[status_q_output, best_questions_output, all_questions_output, formatted_quiz_output]
        )

        best_questions_output.change(
            format_quiz_for_ui,
            inputs=[best_questions_output],
            outputs=[formatted_quiz_output]
        )

        edit_button.click(
            load_quiz_for_editing,
            inputs=[formatted_quiz_output],
            outputs=[edit_status, question_editor, questions_state, index_state, edited_state, next_button]
        )

        load_file_button.click(
            load_file_for_editing,
            inputs=[file_upload],
            outputs=[edit_status, question_editor, questions_state, index_state, edited_state, next_button]
        )

        next_button.click(
            accept_and_next,
            inputs=[question_editor, questions_state, index_state, edited_state],
            outputs=[edit_status, question_editor, questions_state, index_state, edited_state, next_button]
        )

        prev_button.click(
            go_previous,
            inputs=[question_editor, questions_state, index_state, edited_state],
            outputs=[edit_status, question_editor, questions_state, index_state, edited_state, next_button]
        )

        download_button.click(
            save_and_download,
            inputs=[question_editor, questions_state, index_state, edited_state],
            outputs=[edit_status, download_file]
        )

    return app

if __name__ == "__main__":
    app = create_ui()
    app.launch()