import gradio as gr
import pandas as pd
import json

def load_data(path: str):
    records = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            try:
                obj = json.loads(line)
                records.append({
                    "prompt": obj.get("Prompt", "").strip(),
                    "model": obj.get("model", "").strip(),
                    "response": obj.get("model response", "").strip()
                })
            except:
                pass
    return pd.DataFrame(records)

# Load once
df = load_data("all_10_responses.jsonl")
all_prompts = sorted(df["prompt"].unique().tolist())
all_models = sorted(df["model"].unique().tolist())

def get_model_responses(chosen_prompt, chosen_models):
    subset = df[
        (df["prompt"] == chosen_prompt) &
        (df["model"].isin(chosen_models))
    ].sort_values("model")
    # Return a list of tuples: [(model_name, response), …]
    return [(row["model"], row["response"]) for _, row in subset.iterrows()]

with gr.Blocks() as demo:
    gr.Markdown("# Model Responses Explorer")

    with gr.Row():
        prompt_dd = gr.Dropdown(all_prompts, label="Select a Prompt")
        models_dd = gr.CheckboxGroup(all_models, value=all_models, label="Select Models")

    with gr.Column():
        output_acc = gr.Accordion(label="Responses", open=False)

    def update_accordion(prompt, models):
        items = get_model_responses(prompt, models)
        # Gradio accordion expects a single gr.Column or gr.Text
        # We’ll return a single Markdown string combining all models
        md = ""
        for m, r in items:
            md += f"### {m}\n\n{r}\n---\n"
        return md

    accordion_text = gr.Markdown()
    prompt_dd.change(fn=update_accordion, inputs=[prompt_dd, models_dd], outputs=accordion_text)
    models_dd.change(fn=update_accordion, inputs=[prompt_dd, models_dd], outputs=accordion_text)

    demo.launch()