import gradio as gr
from pixaris.feedback_handlers.base import FeedbackHandler

PROJECTS = []


def render_feedback_tab(
    feedback_handler: FeedbackHandler,
):
    # initially load all projects
    global PROJECTS
    PROJECTS = [""] + feedback_handler.load_projects_list()

    feedback_details = (
        gr.State(  # is adjusted from inside a gr.render decorated function. See below.
            value={
                "project": "",
                "feedback_iteration": "",
                "image_name": "",
                "feedback_indicator": False,
                "comment": "",
            }
        )
    )

    def adjust_feedback_details(
        img_path: str, feedback_indicator: bool, comment: str, previous_details: dict
    ):
        """
        This function is only here to adjust the feedback details (that are a gr.State object). This is
        necessary because the feedback details are determined within a function that has the gr.render
        decorator. Thus, the only way to trigger a function within the gr.render function (in this case
        we want to write the feedback) is to write the inputs to a gr.State object and then call the
        respective function within here.
        See here https://www.gradio.app/guides/state-in-blocks (see example with cart).
        """
        previous_details["project"] = img_path.split("/")[-4]  # todo make more robust
        previous_details["feedback_iteration"] = img_path.split("/")[-2]
        previous_details["image_name"] = img_path.split("/")[-1]
        previous_details["feedback_indicator"] = feedback_indicator
        previous_details["comment"] = comment
        feedback_handler.write_single_feedback(feedback=previous_details)
        return previous_details

    with gr.Row():
        with gr.Column(scale=2, min_width=250):
            project_name = gr.Dropdown(
                choices=PROJECTS,
                label="Project",
                filterable=True,
            )

            # initialise hidden feedback iterations and define updating function
            feedback_iterations = gr.Dropdown(visible=False)

            def update_feedback_iteration_choices(project_name, feedback_iterations):
                """Update choices of feedback iterations for selected project."""
                feedback_handler.load_all_feedback_iteration_info_for_project(
                    project_name
                )
                feedback_iteration_choices = feedback_handler.feedback_iteration_choices

                feedback_iterations = gr.Dropdown(
                    label="Feedback Iterations",
                    value=None,
                    choices=feedback_iteration_choices,
                    visible=True,
                    filterable=True,
                    multiselect=True,
                    max_choices=100,
                    interactive=True,
                )
                return feedback_iterations

            # update feedback iterations choices upon project change
            project_name.change(
                fn=update_feedback_iteration_choices,
                inputs=[
                    project_name,
                    feedback_iterations,
                ],
                outputs=[feedback_iterations],
            )

            reload_projects_button = gr.Button(
                "Reload projects",
                variant="secondary",
                interactive=True,
                size="sm",
            )

            def reload_projects(project_name):
                global PROJECTS
                PROJECTS = feedback_handler.load_projects_list()
                project_name = gr.Dropdown(
                    choices=PROJECTS,
                    label="Project",
                    filterable=True,
                )
                return project_name

            reload_projects_button.click(
                fn=reload_projects,
                inputs=[project_name],
                outputs=[project_name],
            )

            columns = gr.Slider(
                minimum=1,
                maximum=7,
                value=5,
                label="Number of images per row",
                step=1,
            )

            display_feedback_checkbox = gr.Checkbox(
                value=False,
                label="Show previous feedback",
                interactive=True,
            )
            sorting_of_images = gr.Radio(
                choices=["image_name", "likes", "dislikes"],
                label="Sorting of images",
                value="image_name",
            )
            reload_feedback_button = gr.Button(
                "Reload feedback",
                variant="secondary",
                interactive=True,
                size="sm",
            )
            reloaded_feedback_count = gr.State(value=1)

            def reload_feedback(project_name, reloaded_feedback_count):
                feedback_handler.load_all_feedback_iteration_info_for_project(
                    project_name
                )
                reloaded_feedback_count += 1
                gr.Info("Reloaded feedback.", duration=2)
                return reloaded_feedback_count

            reload_feedback_button.click(
                fn=reload_feedback,
                inputs=[project_name, reloaded_feedback_count],
                outputs=[reloaded_feedback_count],
            )

        # Main content on the right
        with gr.Column(scale=8):

            @gr.render(
                inputs=[
                    feedback_iterations,
                    columns,
                    display_feedback_checkbox,
                    sorting_of_images,
                    reloaded_feedback_count,
                ]
            )
            def render_images_per_iteration(
                feedback_iterations,
                columns,
                display_feedback_checkbox,
                sorting_of_images,
                reloaded_feedback_count,
            ):
                """
                This function renders the images for each feedback iteration. It is decorated with gr.render
                to allow for dynamic rendering of the images based on the selected feedback iterations.
                - for each feedback_iteration, there is a separate accordion
                - each accordion contains rows of images with the number of columns specified by the user
                - each image is associated with feedback functionality
                """
                if not feedback_iterations:
                    gr.Markdown("No feedback iteration selected.")
                    return
                for feedback_iteration in feedback_iterations:
                    # load the images corresponding to this feedback iteration
                    feedback_iteration_images = (
                        feedback_handler.load_images_for_feedback_iteration(
                            feedback_iteration, sorting_of_images
                        )
                    )

                    # split images into batches of number of columns
                    num_images = len(feedback_iteration_images)
                    images_batches = [
                        feedback_iteration_images[i : i + columns]
                        for i in range(0, num_images, columns)
                    ]

                    # fill up last batch with Nones to have same number of images in each row
                    if len(images_batches[-1]) < columns:
                        images_batches[-1] += [None] * (
                            columns - len(images_batches[-1])
                        )

                    # render images
                    min_width_elements = "10px"
                    with gr.Accordion(label=f"Iteration {feedback_iteration}"):
                        for batch in images_batches:
                            with gr.Row(variant="compact"):
                                for img in batch:
                                    # string of image_name is needed later on, img will be modified by gradio hereafter.
                                    img_path = str(img)
                                    img_name = img_path.split("/")[-1]
                                    # only display image with buttons if it exists
                                    element_visible_bool = bool(img)
                                    with gr.Column(
                                        variant="compact", min_width=min_width_elements
                                    ):
                                        gr.Image(
                                            value=img,
                                            label=img_name,
                                            show_download_button=True,
                                            show_fullscreen_button=True,
                                            visible=element_visible_bool,
                                            min_width=min_width_elements,
                                            scale=1,
                                        )
                                        img_textbox = gr.Textbox(
                                            value=img_path, visible=False
                                        )  # needed bc gr.render
                                        comment = gr.Textbox(
                                            label="Comment",
                                            value="",
                                            visible=element_visible_bool,
                                            min_width=min_width_elements,
                                            scale=1,
                                            interactive=True,
                                        )
                                        feedback_indicator = gr.Radio(
                                            choices=["Like", "Dislike"],
                                            label="Rating",
                                            visible=element_visible_bool,
                                        )
                                        feedback_button = gr.Button(
                                            "Save this image feedback",
                                            visible=element_visible_bool,
                                            size="sm",
                                            min_width=min_width_elements,
                                            scale=1,
                                            variant="primary",
                                            interactive=False,
                                        )

                                        # feedback button is only clickable if feedback indicator is changed recently
                                        def change_feedback_button_interactivity(
                                            feedback_indicator,
                                            feedback_button,
                                        ):
                                            if feedback_indicator:
                                                feedback_button = gr.Button(
                                                    interactive=True
                                                )
                                            else:
                                                feedback_button = gr.Button(
                                                    interactive=False
                                                )

                                            return feedback_button

                                        feedback_indicator.change(
                                            fn=change_feedback_button_interactivity,
                                            inputs=[
                                                feedback_indicator,
                                                feedback_button,
                                            ],
                                            outputs=[feedback_button],
                                        )

                                        # adjusts gr.State object with feedback details and writes feedback
                                        feedback_button.click(
                                            fn=adjust_feedback_details,
                                            inputs=[
                                                img_textbox,
                                                feedback_indicator,
                                                comment,
                                                feedback_details,
                                            ],
                                            outputs=[feedback_details],
                                        ).then(
                                            # change feedback button to non-interactive again (avoid double feedback)
                                            fn=lambda: gr.update(interactive=False),
                                            inputs=None,
                                            outputs=feedback_button,
                                        )

                                        feedback = (
                                            feedback_handler.get_feedback_per_image(
                                                feedback_iteration=feedback_iteration,
                                                image_name=img_name,
                                            )
                                        )
                                        gr.Markdown(
                                            label="Previous Feedback",
                                            value=f"Likes: {feedback['likes']} - Comments: {feedback['comments_liked']}"
                                            if feedback and reloaded_feedback_count > 0
                                            else "",
                                            visible=display_feedback_checkbox
                                            and element_visible_bool,
                                        )
                                        gr.Markdown(
                                            label="Previous Feedback",
                                            value=f"Dislikes: {feedback['dislikes']} - Comments: {feedback['comments_disliked']}"
                                            if feedback and reloaded_feedback_count > 0
                                            else "",
                                            visible=display_feedback_checkbox
                                            and element_visible_bool,
                                        )
