import streamlit as st
import json
import os

# 设置JSON文件的目录
JSON_DIR = 'table_result'

# 加载指定的JSON文件
def load_json(file_name):
    file_path = os.path.join(JSON_DIR, file_name)
    with open(file_path, "r", encoding="utf-8") as f:
        return json.load(f)

# 显示单个字典的信息
# def display_dict(data):
#     st.write("### 文件信息")
#     st.write(f"**Path:** {data['path']}")
#     st.write(f"**Table ID:** {data['table_id']}")
#     st.write(f"**Section:** {data['section']}")
#     st.write("### Table")
#     st.markdown(data['table_html'], unsafe_allow_html=True)

#     st.write("### Context")
#     # 拼接 all_context 并高亮 target_context_ids 的句子
#     all_context = data["all_context"]
#     highlighted_context = ""
#     for idx, sentence in enumerate(all_context):
#         if idx == data["perturb_sentence_id"]:
#             highlighted_context += f"<span style='color:red;'>{sentence}</span> "
#         elif idx in data["target_context_ids"]:
#             highlighted_context += f"**{sentence}** "
#         else:
#             highlighted_context += sentence + " "
#     st.markdown(highlighted_context, unsafe_allow_html=True)

#     st.write("### Selected Paragraphs")
#     for paragraph in data["selected_paragraphs"]:
#         st.write(paragraph)

#     st.write("### Output")
#     st.write("**Perturbed Statement:**")
#     st.write(data["output"]["perturbed_statement"])
#     st.write("**Perturbed Explanation:**")
#     st.write(data["output"]["perturbed_explanation"])
# 修改 Context 渲染逻辑，添加 CSS 控制布局
def display_dict(data):
    st.write("### 文件信息")
    st.write(f"**Path:** {data['path']} | **Table ID:** {data['table_id']} | **Section:** {data['section']}")

    # 使用两个列来分隔内容
    col1, col2 = st.columns(2)

    with col1:
        st.write("### Table")
        st.markdown(data['table_html'], unsafe_allow_html=True)

    with col2:
        st.write("### Context")
        # 拼接 all_context，并高亮特定部分
        all_context = data["all_context"]
        highlighted_context = '<div class="context">'
        for idx, sentence in enumerate(all_context):
            if idx == data["perturb_sentence_id"]:
                highlighted_context += f"<span style='color:red;'>{sentence}</span> "
            elif idx in data["target_context_ids"]:
                highlighted_context += f"<b>{sentence}</b> "
            else:
                highlighted_context += sentence + " "
        highlighted_context += '</div>'
        st.markdown(highlighted_context, unsafe_allow_html=True)

    # 单独分离的内容
    st.write("### Selected Paragraphs")
    for paragraph in data["selected_paragraphs"]:
        st.write(paragraph)

    st.write("### Output")
    st.write("**Perturbed Statement:**")
    st.write(data["output"]["perturbed_statement"])
    st.write("**Perturbed Explanation:**")
    st.write(data["output"]["perturbed_explanation"])


# 主程序
def main():
    st.title("Perturb For Table")

    # 获取文件列表
    file_list = [f for f in os.listdir(JSON_DIR) if f.endswith(".json")]

    # 搜索框选择文件
    selected_file = st.selectbox("选择一个文件", file_list)

    # 加载并展示JSON内容
    if selected_file:
        st.write(f"当前选择文件: **{selected_file}**")
        json_data = load_json(selected_file)

        # 展示每个字典的信息
        for idx, single_dict in enumerate(json_data):
            st.write(f"## 数据 {idx + 1}")
            display_dict(single_dict)

if __name__ == "__main__":
    main()