import gradio as gr
from chat import get_chat_and_score_df, update_chat_display

def create_exploration_tab(df, MODELS, DATASETS, SCORES, HEADER_CONTENT):

    def filter_and_update_display(model, dataset, min_score, max_score, current_index):
        try:
            df_chat = get_chat_and_score_df(model, dataset)

            # Filter by score range
            df_chat = df_chat[
                (df_chat["score"] >= min_score) & (df_chat["score"] <= max_score)
            ]

            if df_chat.empty:
                return (
                    "<div>No data available for selected filters</div>",
                    "<div>No metrics available</div>",
                    "<div>No tool information available</div>",
                    "0/0",
                )

            max_index = len(df_chat) - 1
            current_index = min(current_index, max_index)
            chat_html, metrics_html, tool_html = update_chat_display(
                df_chat, current_index
            )

            return (
                chat_html,
                metrics_html,
                tool_html,
                f"{current_index + 1}/{len(df_chat)}",
            )
        except Exception as e:
            print(f"Error in filter_and_update_display: {str(e)}")
            return (
                f"<div>Error: {str(e)}</div>",
                "<div>No metrics available</div>",
                "<div>No tool information available</div>",
                "0/0",
            )

    with gr.Tab("Data Exploration"):
        gr.HTML(HEADER_CONTENT)

        # All filters in a single row with consistent sizing
        with gr.Row(equal_height=True):
            explore_model = gr.Dropdown(
                choices=MODELS,
                value=MODELS[0],
                label="Model",
                container=True,
                scale=1,
            )
            explore_dataset = gr.Dropdown(
                choices=DATASETS,
                value=DATASETS[0],
                label="Dataset",
                container=True,
                scale=1,
            )
            min_score = gr.Slider(
                minimum=min(SCORES),
                maximum=max(SCORES),
                value=min(SCORES),
                step=0.1,
                label="Minimum Score - TSQ",
                container=True,
                scale=1,
            )
            max_score = gr.Slider(
                minimum=min(SCORES),
                maximum=max(SCORES),
                value=max(SCORES),
                step=0.1,
                label="Maximum Score - TSQ",
                container=True,
                scale=1,
            )

        # Navigation row
        with gr.Row(variant="panel"):
            index_display = gr.HTML(  # Changed the variable name to index_display
                value="0/0", elem_id="index-display", elem_classes="text-center"
            )
            with gr.Row():
                prev_btn = gr.Button("← Previous", size="lg", variant="secondary")
                next_btn = gr.Button("Next →", size="lg", variant="secondary")

        # Content area with equal column widths
        with gr.Row(equal_height=True):
            chat_display = gr.HTML()
            metrics_display = gr.HTML()
            tool_info_display = gr.HTML()

        current_index = gr.State(value=0)

        # Update display on filter change
        def update_on_filter_change(model, dataset, min_score, max_score):
            return filter_and_update_display(model, dataset, min_score, max_score, 0)

        for control in [explore_model, explore_dataset, min_score, max_score]:
            control.change(
                update_on_filter_change,
                inputs=[explore_model, explore_dataset, min_score, max_score],
                outputs=[
                    chat_display,
                    metrics_display,
                    tool_info_display,
                    index_display,
                ],  # Changed to index_display
            )

        # Navigation functions
        def navigate(direction, current_idx, model, dataset, min_score, max_score):
            new_index = current_idx + direction
            return (
                *filter_and_update_display(
                    model, dataset, min_score, max_score, new_index
                ),
                new_index,
            )

        prev_btn.click(
            lambda idx, m, d, min_s, max_s: navigate(-1, idx, m, d, min_s, max_s),
            inputs=[
                current_index,
                explore_model,
                explore_dataset,
                min_score,
                max_score,
            ],
            outputs=[
                chat_display,
                metrics_display,
                tool_info_display,
                index_display,
                current_index,
            ],  # Changed to index_display
        )

        next_btn.click(
            lambda idx, m, d, min_s, max_s: navigate(1, idx, m, d, min_s, max_s),
            inputs=[
                current_index,
                explore_model,
                explore_dataset,
                min_score,
                max_score,
            ],
            outputs=[
                chat_display,
                metrics_display,
                tool_info_display,
                index_display,
                current_index,
            ],  # Changed to index_display
        )

        return (
            chat_display,
            metrics_display,
            tool_info_display,
            index_display,  # Changed to index_display
        )