import gradio as gr
from global_var import get_datasets
from tinydb import Query

def dataset_manage_page():
    with gr.Blocks() as demo:
        gr.Markdown("## 数据集管理")
        with gr.Row():
            # 获取数据集列表并设置初始值
            datasets_list = [str(ds["name"]) for ds in get_datasets().all()]
            initial_dataset = datasets_list[0] if datasets_list else None
            
            dataset_dropdown = gr.Dropdown(
                choices=datasets_list,
                value=initial_dataset,  # 设置初始选中项
                label="选择数据集",
                allow_custom_value=True,
                interactive=True
            )

        # 添加数据集展示组件
        qa_dataset = gr.Dataset(
            components=["text", "text"],
            label="问答数据",
            headers=["问题", "答案"],
            samples=[["示例问题", "示例答案"]],
            samples_per_page=20,
        )

        def update_qa_display(dataset_name):
            if not dataset_name:
                return {"samples": [], "__type__": "update"}
            
            # 从数据库获取数据集
            Dataset = Query()
            ds = get_datasets().get(Dataset.name == dataset_name)
            if not ds:
                return {"samples": [], "__type__": "update"}
            
            # 提取所有Q_A数据
            qa_list = []
            for item in ds["dataset_items"]:
                for qa in item["message"]:
                    qa_list.append([qa["question"], qa["answer"]])
            
            return {"samples": qa_list, "__type__": "update"}

        # 绑定事件，更新QA数据显示
        dataset_dropdown.change(
            update_qa_display,
            inputs=dataset_dropdown,
            outputs=qa_dataset
        )

    return demo