#!/usr/bin/env python3
"""
简单的反馈收集工具
在MeanAudio生成音频后，运行此脚本收集用户偏好
"""

import json
import os
import sys
from datetime import datetime
from pathlib import Path
import gradio as gr

# 设置反馈目录
FEEDBACK_DIR = Path("./rlhf_feedback")
FEEDBACK_DIR.mkdir(exist_ok=True)
FEEDBACK_FILE = FEEDBACK_DIR / "user_preferences.jsonl"

def save_feedback(audio1_path, audio2_path, prompt, preference, comment=""):
    """保存反馈数据"""
    feedback_data = {
        "timestamp": datetime.now().isoformat(),
        "prompt": prompt,
        "audio1_path": audio1_path,
        "audio2_path": audio2_path,
        "preference": preference,
        "additional_comment": comment
    }
    
    with open(FEEDBACK_FILE, "a", encoding="utf-8") as f:
        f.write(json.dumps(feedback_data, ensure_ascii=False) + "\n")
    
    return f"✅ 反馈已保存！偏好: {preference}"

def create_feedback_interface():
    """创建反馈收集界面"""
    
    with gr.Blocks(title="MeanAudio 反馈收集器") as demo:
        gr.Markdown("# MeanAudio 反馈收集器")
        gr.Markdown("*请输入生成的音频文件路径和提示词，然后选择您的偏好*")
        
        with gr.Row():
            with gr.Column():
                prompt_input = gr.Textbox(
                    label="提示词",
                    placeholder="输入用于生成音频的提示词..."
                )
                
                audio1_path = gr.Textbox(
                    label="音频文件1路径",
                    placeholder="./output/gradio/prompt_timestamp_0.flac"
                )
                
                audio2_path = gr.Textbox(
                    label="音频文件2路径", 
                    placeholder="./output/gradio/prompt_timestamp_1.flac"
                )
                
            with gr.Column():
                # 显示音频
                audio1_player = gr.Audio(label="音频1")
                audio2_player = gr.Audio(label="音频2")
                
                load_btn = gr.Button("🔄 加载音频文件")
        
        # 反馈区域
        gr.Markdown("---")
        gr.Markdown("### 请选择您的偏好")
        
        preference = gr.Radio(
            choices=[
                ("音频1更好", "audio1"),
                ("音频2更好", "audio2"),
                ("两者质量相等", "equal"),
                ("两者都不好", "both_bad")
            ],
            label="哪个音频更好？"
        )
        
        comment = gr.Textbox(
            label="额外评论 (可选)",
            placeholder="关于音频质量的具体反馈...",
            lines=3
        )
        
        submit_btn = gr.Button("📝 提交反馈", variant="primary")
        
        result = gr.Textbox(label="结果", interactive=False)
        
        # 事件处理
        def load_audio_files(path1, path2):
            """加载音频文件用于播放"""
            audio1 = path1 if os.path.exists(path1) else None
            audio2 = path2 if os.path.exists(path2) else None
            return audio1, audio2
        
        load_btn.click(
            fn=load_audio_files,
            inputs=[audio1_path, audio2_path],
            outputs=[audio1_player, audio2_player]
        )
        
        submit_btn.click(
            fn=save_feedback,
            inputs=[audio1_path, audio2_path, prompt_input, preference, comment],
            outputs=[result]
        )
        
        # 使用说明
        gr.Markdown("---")
        gr.Markdown("""
        ### 使用说明
        1. 先运行 MeanAudio 生成两个音频文件
        2. 将生成的音频文件路径复制到上面的输入框中
        3. 点击"加载音频文件"来播放音频
        4. 选择您的偏好并提交反馈
        5. 反馈数据将保存到 `./rlhf_feedback/user_preferences.jsonl`
        6. 使用 `python analyze_feedback.py` 分析收集的反馈数据
        """)
    
    return demo

if __name__ == "__main__":
    demo = create_feedback_interface()
    print("启动反馈收集界面...")
    print(f"反馈数据将保存到: {FEEDBACK_FILE}")
    demo.launch(server_name="127.0.0.1", server_port=7861)