# import gradio as gr
# from data_handler import get_next_sample, annotate_text, get_generation_columns, load_progress


import gradio as gr
from data_handler import get_next_sample, annotate_text, get_generation_columns, load_progress

# Track the progress of all annotations
progress = load_progress()  # Load previous progress from the progress file
annotation_history = list(progress.keys())  # Maintain a list of comment IDs in order of annotation
current_annotation_idx = -1  # Start from no previous annotation

def update_ui(dataset, generation_columns):
    """Update UI with the next sample or show completion message."""
    sample = get_next_sample(dataset, progress)

    if sample is None:
        return gr.update(visible=False), gr.update(value="All comments annotated!"), gr.update(visible=False)

    # Ensure to return the correct number of outputs for Gradio (textboxes + comment_id)
    textboxes = [gr.update(value=sample[col]) for col in generation_columns]
    return textboxes + [gr.update(value=sample["comment_id"])]  # Return sample comment_id

def go_back_ui():
    """Go back to the previous annotated sample in the history."""
    global current_annotation_idx

    # Check if we have any previous annotations
    if current_annotation_idx <= 0:
        return gr.update(visible=False), gr.update(value="No previous annotations.")

    # Move one step back
    current_annotation_idx -= 1

    # Get the previous annotated sample based on the current_annotation_idx
    prev_annotated_id = annotation_history[current_annotation_idx]
    prev_sample = progress[prev_annotated_id]

    textboxes = [
        gr.update(value=prev_sample["grammar"]),
        gr.update(value=prev_sample["coherence"]),
        gr.update(value=prev_sample["preferred_text"])
    ]
    return textboxes + [gr.update(value=prev_annotated_id)]  # Show the previous annotation

def create_ui(dataset, current_index):
    """Creates the Gradio UI for annotation."""
    generation_columns = get_generation_columns(dataset)

    with gr.Blocks() as demo:
        gr.Markdown("## Annotate Generated Texts")
        gr.Markdown("#### Rate each generation and select the best one")

        textboxes, grammar_radios, coherence_radios = [], [], []

        for col in generation_columns:
            with gr.Row():
                with gr.Column(scale=3):
                    gr.Markdown(f"### {col}")
                    text_input = gr.Textbox(label=f"Generated Text ({col})", lines=5)
                    textboxes.append(text_input)

                with gr.Column(scale=1):
                    grammar = gr.Radio(choices=["-2", "-1", "0", "1", "2"], label="Grammar", interactive=True)
                    coherence = gr.Radio(choices=["-2", "-1", "0", "1", "2"], label="Coherence", interactive=True)
                    grammar_radios.append(grammar)
                    coherence_radios.append(coherence)

        preferred_text = gr.Dropdown(choices=generation_columns, label="Select Best Generation")
        comment_id_box = gr.Textbox(label="Comment ID", interactive=False)  # Store comment_id

        with gr.Row():
            submit_btn = gr.Button("Submit Annotation")
            go_back_btn = gr.Button("Go Back to Last Annotation")

        output_message = gr.Textbox(label="Status", interactive=False)

        # Set initial values
        demo.load(lambda: update_ui(dataset, generation_columns), outputs=textboxes + [comment_id_box])

        # Define button actions
        submit_btn.click(
            lambda grammar_scores, coherence_scores, preferred, comment_id: annotate_text(
                dataset, comment_id, grammar_scores, coherence_scores, preferred, generation_columns, progress
            ),
            inputs=grammar_radios + coherence_radios + [preferred_text, comment_id_box],
            outputs=textboxes + [output_message]
        )

        go_back_btn.click(
            go_back_ui,
            outputs=textboxes + [comment_id_box, output_message]
        )

    return demo


# # Track the progress of all annotations
# progress = load_progress()  # Load previous progress from the progress file
# annotation_history = list(progress.keys())  # Maintain a list of comment IDs in order of annotation
# current_annotation_idx = -1  # Start from no previous annotation

# def update_ui(dataset, generation_columns):
#     """Update UI with the next sample or show completion message."""
#     sample = get_next_sample(dataset, progress)

#     if sample is None:
#         return gr.update(visible=False), gr.update(value="All comments annotated!"), gr.update(visible=False)

#     textboxes = [gr.update(value=sample[col]) for col in generation_columns]
#     return textboxes + [gr.update(value=sample["comment_id"])]  # Return sample comment_id

# def go_back_ui():
#     """Go back to the previous annotated sample in the history."""
#     global current_annotation_idx

#     # Check if we have any previous annotations
#     if current_annotation_idx <= 0:
#         return gr.update(visible=False), gr.update(value="No previous annotations.")

#     # Move one step back
#     current_annotation_idx -= 1

#     # Get the previous annotated sample based on the current_annotation_idx
#     prev_annotated_id = annotation_history[current_annotation_idx]
#     prev_sample = progress[prev_annotated_id]

#     textboxes = [gr.update(value=prev_sample["grammar"]),
#                  gr.update(value=prev_sample["coherence"]),
#                  gr.update(value=prev_sample["preferred_text"])]

#     return textboxes + [gr.update(value=prev_annotated_id)]  # Show the previous annotation

# def create_ui(dataset, current_index):
#     """Creates the Gradio UI for annotation."""
#     generation_columns = get_generation_columns(dataset)

#     with gr.Blocks() as demo:
#         gr.Markdown("## Annotate Generated Texts")
#         gr.Markdown("#### Rate each generation and select the best one")

#         textboxes, grammar_radios, coherence_radios = [], [], []

#         for col in generation_columns:
#             with gr.Row():
#                 with gr.Column(scale=3):
#                     gr.Markdown(f"### {col}")
#                     text_input = gr.Textbox(label=f"Generated Text ({col})", lines=5)
#                     textboxes.append(text_input)

#                 with gr.Column(scale=1):
#                     grammar = gr.Radio(choices=["-2", "-1", "0", "1", "2"], label="Grammar", interactive=True)
#                     coherence = gr.Radio(choices=["-2", "-1", "0", "1", "2"], label="Coherence", interactive=True)
#                     grammar_radios.append(grammar)
#                     coherence_radios.append(coherence)

#         preferred_text = gr.Dropdown(choices=generation_columns, label="Select Best Generation")
#         comment_id_box = gr.Textbox(label="Comment ID", interactive=False)  # Store comment_id

#         with gr.Row():
#             submit_btn = gr.Button("Submit Annotation")
#             go_back_btn = gr.Button("Go Back to Last Annotation")

#         output_message = gr.Textbox(label="Status", interactive=False)

#         # Set initial values
#         demo.load(lambda: update_ui(dataset, generation_columns), outputs=textboxes + [comment_id_box])

#         # Define button actions
#         submit_btn.click(
#             lambda grammar_scores, coherence_scores, preferred, comment_id: annotate_text(
#                 dataset, comment_id, grammar_scores, coherence_scores, preferred, generation_columns, progress
#             ),
#             inputs=grammar_radios + coherence_radios + [preferred_text, comment_id_box],
#             outputs=textboxes + [output_message]
#         )

#         go_back_btn.click(
#             go_back_ui,
#             outputs=textboxes + [comment_id_box, output_message]
#         )

#     return demo


# # # Store the history of annotated samples
# # last_sample = None  # To store the previous sample for "Go Back"

# # def update_ui(dataset, generation_columns):
# #     """Update UI with the next sample or show completion message."""
# #     global last_sample
# #     sample = get_next_sample(dataset)

# #     if sample is None:
# #         return gr.update(visible=False), gr.update(value="All comments annotated!"), gr.update(visible=False)

# #     last_sample = sample  # Store the current sample for going back
# #     textboxes = [gr.update(value=sample[col]) for col in generation_columns]
    
# #     return textboxes + [gr.update(value=sample["comment_id"])]  # Return sample comment_id

# # def go_back_ui(dataset, generation_columns):
# #     """Return the previous annotated sample (if available)."""
# #     global last_sample
# #     if last_sample is None:
# #         return gr.update(visible=False), gr.update(value="No previous sample.")

# #     textboxes = [gr.update(value=last_sample[col]) for col in generation_columns]
    
# #     return textboxes + [gr.update(value=last_sample["comment_id"])]  # Display previous comment_id

# # def create_ui(dataset):
# #     """Creates the Gradio UI for annotation."""
# #     generation_columns = get_generation_columns(dataset)

# #     with gr.Blocks() as demo:
# #         gr.Markdown("## Annotate Generated Texts")
# #         gr.Markdown("#### Rate each generation and select the best one")

# #         textboxes, grammar_radios, coherence_radios = [], [], []

# #         for col in generation_columns:
# #             with gr.Row():
# #                 with gr.Column(scale=3):
# #                     gr.Markdown(f"### {col}")
# #                     text_input = gr.Textbox(label=f"Generated Text ({col})", lines=5)
# #                     textboxes.append(text_input)

# #                 with gr.Column(scale=1):
# #                     grammar = gr.Radio(choices=["-2", "-1", "0", "1", "2"], label="Grammar", interactive=True)
# #                     coherence = gr.Radio(choices=["-2", "-1", "0", "1", "2"], label="Coherence", interactive=True)
# #                     grammar_radios.append(grammar)
# #                     coherence_radios.append(coherence)

# #         preferred_text = gr.Dropdown(choices=generation_columns, label="Select Best Generation")
# #         comment_id_box = gr.Textbox(label="Comment ID", interactive=False)  # Store comment_id

# #         with gr.Row():
# #             submit_btn = gr.Button("Submit Annotation")
# #             go_back_btn = gr.Button("Go Back to Last Annotation")

# #         output_message = gr.Textbox(label="Status", interactive=False)

# #         # Set initial values
# #         demo.load(lambda: update_ui(dataset, generation_columns), outputs=textboxes + [comment_id_box])

# #         # Define button actions
# #         submit_btn.click(
# #             lambda grammar_scores, coherence_scores, preferred, comment_id: annotate_text(
# #                 dataset, comment_id, grammar_scores, coherence_scores, preferred, generation_columns
# #             ),
# #             inputs=grammar_radios + coherence_radios + [preferred_text, comment_id_box],
# #             outputs=textboxes + [output_message]
# #         )

# #         go_back_btn.click(
# #             lambda: go_back_ui(dataset, generation_columns),
# #             outputs=textboxes + [comment_id_box, output_message]
# #         )

# #     return demo