import gradio as gr
import config
from ui_components import create_ui
from annotation import load_annotations
from task_manager import get_current_task, get_current_task_with_annotations, get_total_tasks

css = """
.center { text-align: center; }
.audio-container { margin: 10px; padding: 15px; }
.selected {
    border: 3px solid #4CAF50 !important;
    background-color: #e8f5e9 !important;
}
.user-task-info {
    font-size: 16px;
    color: #333;
    padding: 10px;
    background-color: #f0f0f0;
    border-radius: 5px;
    display: flex;
    justify-content: space-between;
    align-items: center;
}
"""


def set_user_info(request: gr.Request):
    """设置用户信息到 State 并加载用户特定的标注"""
    username = request.username if hasattr(request, 'username') else "unknown"
    annotation_results = load_annotations(username)
    print(f"加载用户 {username} 的标注结果：{annotation_results}")

    # 获取当前任务信息，并考虑已有的标注结果
    inst, text, audioA_update, audioB_update, prev_disabled, next_disabled, task_num = get_current_task_with_annotations(
        annotation_results)

    # 获取总任务数
    total_tasks = get_total_tasks()

    # 合并用户信息和任务编号到任务编号位置
    combined_task_info = f'<div class="user-task-info"><span>👤 当前用户: {username}</span><span><strong>任务编号: {task_num} / {total_tasks}</strong></span></div>'

    # 返回所有需要初始化的组件值
    return (
        username,
        annotation_results,
        inst,  # instruction
        text,  # text_box
        audioA_update,  # audioA
        audioB_update,  # audioB
        gr.update(interactive=not prev_disabled),  # btn_prev
        gr.update(interactive=not next_disabled),  # btn_next
        gr.update(value=combined_task_info)  # task_number 位置显示合并信息
    )


if __name__ == "__main__":
    print("启动应用...")
    with gr.Blocks(css=css) as demo:
        username = gr.State(value="unknown")
        annotation_results = gr.State(value={})
        ui_components = create_ui(get_current_task(), username, annotation_results)

        # 修改 demo.load，移除 user_display，保留 task_number
        demo.load(
            set_user_info,
            inputs=None,
            outputs=[
                username,
                annotation_results,
                ui_components["instruction"],
                ui_components["text_box"],
                ui_components["audioA"],
                ui_components["audioB"],
                ui_components["btn_prev"],
                ui_components["btn_next"],
                ui_components["task_number"]  # 现在包含用户信息和任务编号
            ]
        )

        # 其余按钮绑定保持不变
        ui_components["btn_A"].click(
            ui_components["select_audio"],
            inputs=[gr.State("A"), ui_components["audioA"], ui_components["audioB"], annotation_results, username],
            outputs=[ui_components["audioA"], ui_components["audioB"], annotation_results]
        )

        ui_components["btn_B"].click(
            ui_components["select_audio"],
            inputs=[gr.State("B"), ui_components["audioA"], ui_components["audioB"], annotation_results, username],
            outputs=[ui_components["audioA"], ui_components["audioB"], annotation_results]
        )

        ui_components["btn_prev"].click(
            ui_components["change_task"],
            inputs=[gr.State("prev"), annotation_results, username],  # 添加 username 参数
            outputs=[ui_components["instruction"], ui_components["text_box"], ui_components["audioA"],
                     ui_components["audioB"], ui_components["btn_prev"], ui_components["btn_next"],
                     ui_components["task_number"], annotation_results]  # 保持 task_number
        )

        ui_components["btn_next"].click(
            ui_components["change_task"],
            inputs=[gr.State("next"), annotation_results, username],  # 添加 username 参数
            outputs=[ui_components["instruction"], ui_components["text_box"], ui_components["audioA"],
                     ui_components["audioB"], ui_components["btn_prev"], ui_components["btn_next"],
                     ui_components["task_number"], annotation_results]  # 保持 task_number
        )

    demo.launch(auth=config.ANNOTATOR, share=True,ssr_mode=False)