import streamlit as st

# 存储标注信息
if "annotations" not in st.session_state:
    st.session_state.annotations = []

# 添加标注


def add_annotation(selected_text, label):
    start_idx = text.find(selected_text)
    if start_idx != -1:
        end_idx = start_idx + len(selected_text)
        st.session_state.annotations.append(
            {"text": selected_text, "label": label, "start": start_idx, "end": end_idx})

# 移除标注


def remove_annotation(annotation):
    st.session_state.annotations = [
        ann for ann in st.session_state.annotations if ann != annotation
    ]

# 生成带按钮的高亮文本


def render_highlighted_text(text, annotations):
    annotated_text = []
    last_idx = 0
    for idx, ann in enumerate(sorted(annotations, key=lambda x: x["start"])):
        # 添加非标注的部分
        annotated_text.append(text[last_idx:ann["start"]])
        # 添加标注的按钮
        if st.button(f"{ann['text']} ({ann['label']})", key=f"annotation-{idx}"):
            st.session_state.clicked_annotation = ann
        last_idx = ann["end"]
    # 添加剩余的部分
    annotated_text.append(text[last_idx:])
    return "".join(annotated_text)


# 主界面
st.title("文本标注工具")
st.write("输入文本并对其中的部分内容进行标注。")

# 输入文本
text = st.text_area("输入文本：", height=200)

# 文本选择与标注
if text:
    st.write("选中文字并进行标注：")
    col1, col2 = st.columns([3, 1])

    # 输入标注的标签
    with col1:
        selected_text = st.text_input("选中的文本：")
    with col2:
        label = st.text_input("标签：")
        if st.button("添加标注"):
            if selected_text and label and selected_text in text:
                add_annotation(selected_text, label)
            else:
                st.warning("请选择有效的文本并输入标签！")

    # 渲染高亮文本
    st.markdown(render_highlighted_text(text, st.session_state.annotations))

    # 点击标注的文本
    if "clicked_annotation" in st.session_state:
        clicked_annotation = st.session_state.clicked_annotation
        st.write(
            f"你点击了: **{clicked_annotation['text']}** (标签: {clicked_annotation['label']})")
        if st.button("删除标注", key="delete"):
            remove_annotation(clicked_annotation)
            del st.session_state.clicked_annotation