import streamlit as st
import os
import runpy

# 从环境变量获取 token
token = os.getenv('token')
weather_token = os.getenv('weather_token')
horoscope_token = os.getenv("horoscope_token")

st.set_page_config(layout="wide", page_title="My Multi-Page App")

def set_env_variable(key, value):
    os.environ[key] = value

def home_page():
    st.header("欢迎来到首页")
    # 检查环境变量是否已存在
    global token, weather_token, horoscope_token
    if not token:
        token = st.text_input("请输入浦语 token:", type="password", key="token")
    if not weather_token:
        weather_token = st.text_input("请输入和风天气 token:", type="password", key="weather_token")
    if not horoscope_token:
        horoscope_token = st.text_input("请输入星座运势 token:", type="password", key="horoscope_token")
    
    if st.button("保存并体验 agent"):
        if token and weather_token and horoscope_token:
            set_env_variable("token", token)  # 设置环境变量
            set_env_variable("weather_token", weather_token)
            set_env_variable("horoscope_token", horoscope_token)
            st.session_state.token_entered = True
            st.rerun()
        else:
            st.error("请输入所有 token")

# 检查会话状态
if 'token_entered' not in st.session_state:
    # 如果 token 存在于环境变量中，直接跳过输入
    st.session_state.token_entered = bool(token and weather_token and horoscope_token)

if not st.session_state.token_entered:
    home_page()
else:
    # 动态加载子页面
    page = st.sidebar.radio("选择页面", ["天气和星座助手", "博客写作助手"])
    if page == "天气和星座助手":
        # 加载天气和星座查询的页面
        runpy.run_path("examples/agent_api_web_demo.py", run_name="__main__")
    elif page == "博客写作助手":
        # 加载博客写作助手的页面
        runpy.run_path("examples/multi_agents_api_web_demo.py", run_name="__main__")