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.markdown('''
        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 and consent to having your conversations used for training language models.**
    ''')

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": "請用繁體中文(zh-tw)回覆，並以台灣社會的場景為主。注意你的中英全半型交雜情況下，要有適當的空格。"}]
    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})