import streamlit as st
import requests
import json
# 定义 generate 函数
def generate(query, state_name):
    url = "http://47.113.113.230:6969/process_text"
    payload = {
        "input_text": query,
        "states_name": state_name
    }
    response = requests.post(url, json=payload)
    if response.status_code == 200:
        result = response.json()
        output_text = json.loads(result["output_text"])
        elapsed_time = result["elapsed_time"]
        st.write(f"生成时间: {elapsed_time} 秒")
        st.write("生成结果:")
        st.json(output_text)
    else:
        st.write("请求失败")

# 定义 states 变量
states = ["unit_extractor", "domain_expert"]

# 设置应用标题
st.title("多States切换生成Demo")

# 初始化 session_state
if 'query' not in st.session_state:
    st.session_state.query = ""
if 'state_name' not in st.session_state:
    st.session_state.state_name = states[0]
# 创建侧边栏
with st.sidebar:
    st.header("Demo Inputs")
    # 定义 demo_inputs 列表
    demo_inputs = [
        {"query": "总共453户", "state": "unit_extractor"},
        {"query": "菌子，又名蘑菇。通常大家把人工种植的叫蘑菇，野生的叫菌子。 菌子是人们常食的山珍佳肴，即食用菌类植物的子实体，指蕈、菇、灵芝一类，多来源于担子菌纲的某些科属。常见的食用菌菌类有香菇、蘑菇、猴头、木耳、灵芝等，多属山珍类，古时均属野生，现在已能人工培植，可以大范围生产。菌子已走出深山老林，既进入了大的酒店餐馆，也进入了寻常百姓家。许多食用真菌营养丰富，为食中佳品，其作用家喻户晓，人人皆知。", "state": "domain_expert"},
        {"query": "总面积约1234.8平方米", "state": "unit_extractor"},
        {"query": "云南省，简称“云”或“滇”，中国23个省之一，位于西南地区，省会昆明。介于北纬21°8′～29°15′，东经97°31′～106°11′之间，东部与贵州、广西为邻，北部与四川相连，西北部紧依西藏，西部与缅甸接壤，南部和老挝、越南毗邻，云南省总面积39.41万平方千米，居全国第8位。云南省是全国边境线最长的省份之一，有8个州（市）的25个边境县分别与缅甸、老挝和越南交界。", "state": "domain_expert"},
        {"query": "包括394.34平方公里", "state": "unit_extractor"},
        {"query": "元谋人，因发现地点在云南元谋县上那蚌村西北小山岗上而得名。元谋人牙齿化石是1965年“五一”节在云南元谋县上那蚌村发现的，元谋县被誉为“元谋人的故乡”。1976年根据古地磁学方法测定，生活年代距今约170万年，差距最多不超过前后十万年（也有学者认为其年代不应超过73万年，即可能为差距在60万至50万年或更晚一些）。", "state": "domain_expert"},
    ]

    # 展示 demo_inputs 列表
    for i, demo in enumerate(demo_inputs):
        st.write(f"Query: {demo['query']}")
        st.write(f"State: {demo['state']}")
        if st.button(f"Apply {i}"):
            st.session_state.query = demo['query']
            st.session_state.state_name = demo['state']
            st.rerun()

# 创建下拉框
state_name = st.selectbox("选择一个State", states, index=states.index(st.session_state.state_name))

# 创建输入框
query = st.text_input("输入查询内容", value=st.session_state.query)

# 创建提交按钮
if st.button("提交"):
    # 调用 generate 函数
    generate(query, state_name)