import streamlit as st
from openai import OpenAI, AzureOpenAI
from typing import Iterator
import os
from phoenix.otel import register
from opencc import OpenCC
cc = OpenCC('s2twp')

PHOENIX_API_KEY=st.secrets['PHOENIX_API_KEY']
os.environ["PHOENIX_CLIENT_HEADERS"] = f"api_key={PHOENIX_API_KEY}"

tracer_provider = register(
  project_name=st.secrets['PHOENIX_PROJECT_NAME'],
  endpoint=os.getenv('PHOENIX_COLLECTOR_ENDPOINT'),
)

from openinference.instrumentation.openai import OpenAIInstrumentor
OpenAIInstrumentor().instrument(tracer_provider=tracer_provider)

st.set_page_config(
    page_title="Free GPT-4o Chat",
    page_icon="💬",
    layout="centered"
)

st.title("💬 GPT-4o Chat")
with st.expander("Notice"):
    st.write('''
        Please note that this app collects your conversation records. 
        Do not input any personal or sensitive information, such as:
        - National identification numbers
        - Contact information (phone numbers, email addresses)
        - Bank account or credit card details
        - Health or medical information
        - Any other data that can be used to identify you

        Use this app cautiously and avoid sharing sensitive data. 
        Do not use this app in inappropriate contexts. By using this app, you agree to these terms and conditions.
    ''')

client = AzureOpenAI(
    api_key=st.secrets['API_KEY'],
    api_version=st.secrets['API_VERSION'],
    azure_endpoint=st.secrets['ENDPOINT']
)

if "openai_model" not in st.session_state:
    st.session_state["openai_model"] = st.secrets['MODEL']

if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

if prompt := st.chat_input("What is up?"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    messages = [{"role": "system", "content": "如果使用者使用繁體中文，請用繁體中文回答使用者，並以台灣社會的場景為主。"}]
    for m in st.session_state.messages:
        messages.append({"role": m["role"], "content": m["content"]})

    with st.chat_message("assistant"):
        stream = client.chat.completions.create(
            model=st.session_state["openai_model"],
            messages=messages,
            stream=True,
        )
        response = st.write_stream(stream)
    st.session_state.messages.append(
        {"role": "assistant", "content": response})