import streamlit as st
import requests

# 创建 Streamlit 页面布局
st.title('填充掩码演示')

# 示例输入数据
demo_inputs = [
    "生活的真谛在于[MASK]。",
    "[MASK]的海滩世界闻名。",
    "雨后，我[MASK]在公园里，呼吸新鲜的空气。",
    "我[MASK]了一部好电影，感觉很愉快。",
    "再次回到[MASK]，我感到很亲切。",
    "一九四九年，新中国[MASK]。",
    "在黑暗的[MASK]里，弱肉强食是唯一的法则。",
]
# 创建页面布局
# col1, col2 = st.columns(2)
#col1 is set to 20% width
#col2 is set to 80% width
col1, col2 = st.columns([2, 6])
with col1:
    st.header("Demo Input")
    for i, input_text in enumerate(demo_inputs):
        # 为每个示例输入创建一个行布局
        row = st.columns([4, 1])  # 调整列的宽度比例
        with row[0]:  # 左边显示文本
            st.write(input_text)
        with row[1]:  # 右边显示选择按钮
            if st.button(f"选择", key=f"btn_{i}"):  # 给每个按钮一个唯一的key
                st.session_state['selected_input'] = input_text

with col2:
    st.header("输入")
    # 如果session_state中有selected_input，则使用它作为输入框的默认值
    user_input = st.text_input("请输入文本", value=st.session_state.get('selected_input', ''))

    # 创建一个按钮，当按下时执行操作
    if st.button('发送请求'):
        # 使用用户输入构造请求URL
        url = f"http://47.113.113.230:6969/fill_mask?input_text={user_input}"
        
        response = requests.post(url)
        
        if response.status_code == 200:
            # 解析响应数据并以表格形式显示结果或错误信息
            response_data = response.json()  # 将响应内容解析为JSON
            # 检查是否有错误信息
            if response_data['msg'] is not None and response_data['msg'] != "":
                # 显示错误信息
                st.error(f"错误: {response_data['msg']}")
            else:
                # 没有错误，准备数据以显示在表格中
                elapsed_time = response_data['elapsed_time']
                # 显示请求的响应时间
                st.text(f"响应时间: {elapsed_time} 秒")
                results = response_data['results']
                # 创建一个空列表用于存储转换后的数据
                table_data = []
                for result in results:
                    # 将每个结果添加到列表中
                    table_data.append({"输出文本": result['output_text'], "得分": result['score']})
                # 使用Streamlit的table函数显示数据
                st.table(table_data)
        else:
            # 如果请求失败，显示错误信息
            st.error('请求失败，请检查输入或服务器状态。')