#import openai
from openai import OpenAI
import streamlit as st
import toml
import base64
import requests

#initial values
if 'chating_now' not in st.session_state: st.session_state['chating_now'] = False
if 'patient_six' not in st.session_state: st.session_state['patient_six'] = ""
if 'patient_age' not in st.session_state: st.session_state['patient_age'] = ""

client = OpenAI(base_url="http://localhost:8888/v1/", api_key="not-needed")

#Steramlit initial -------------------------------------------------------------------
secrets = toml.load("streamlit/secrets.toml")

st.set_page_config(layout="wide")
st.write('<style>div.row-widget.stRadio > div{flex-direction:row;}</style>', unsafe_allow_html=True)
#將頂端的空白減少
st.markdown(
    """
        <style>
            .appview-container .main .block-container {{
                padding-top: {padding_top}rem;
                padding-bottom: {padding_bottom}rem;
                }}

        </style>""".format(
        padding_top=1, padding_bottom=1
    ),
    unsafe_allow_html=True,
)
colA, colB = st.columns([0.3,2])
with colA:    
    st.image('doctor.png', width=100)
with colB:
    st.title("TAIDE 7B - 醫療AI問答系統")

# Store the initial value of widgets in session state
if "visibility" not in st.session_state:
    st.session_state.visibility = "visible"
    st.session_state.disabled = False
if "visibility" not in st.session_state:
    st.session_state.moredesc = False

if "exec_moredesc" not in st.session_state:
    st.session_state.exec_moredesc = False    
#-----------------------------------------------------------------------------------------

last_prompt = ""
last_response = ""
q_input = ""
form_login = st.empty()
form_logout = st.empty()

with form_login.container():
   
    with st.form(key='patient_form'):
        col1, col2, col3 = st.columns([0.5,0.5,0.5])
        with col1:
            patient_six = st.radio(
                "性別",
                ["男性", "女性"],
                index=None,
                label_visibility = "collapsed"
            )
        with col2:
            patient_age = st.selectbox(
                '年齡區間',
                ('0-9歲', '10-19歲', '20-29歲', '30-39歲', '40-49歲', '50-59歲', '60-69歲', '70-79歲', '80-89歲', '90-99歲', '100歲以上'),
                index=None, label_visibility = "collapsed")
        with col3:
            submit = st.form_submit_button(label='開始咨詢',use_container_width = True, type='primary')
            
        if submit:
            st.session_state.patient_six = patient_six
            st.session_state.patient_age = patient_age 
            st.session_state.chating_now = True

with form_logout.container():
    col1, col2, col3 = st.columns([0.5,0.5,0.5])
    with col1:
        if "男" in st.session_state.patient_six:
            sixtxt = ":blue[男性] :boy:"
        else:
            sixtxt = ":blue[女性] :girl:"            

        st.markdown("您的性別: {}".format(sixtxt))

    with col2:    
        st.write("您的年齡區間: {}".format(st.session_state.patient_age))

    with col3:  
        if st.button("結束咨詢", key="end",use_container_width = True, type='primary'):
            st.session_state.patient_six = ""
            st.session_state.patient_age = ""
            st.session_state.chating_now = False
            st.session_state.messages = []


if st.session_state.patient_six=="" or st.session_state.patient_age=="" or st.session_state.chating_now is False:        
    st.session_state.patient_six = ""
    st.session_state.patient_age = "" 
    st.session_state.chating_now = False
    form_logout.empty()

else:
    form_login.empty()

#st.write(st.session_state.patient_six)
#st.write(st.session_state.patient_age)
#st.write(st.session_state.chating_now)      
if st.session_state.chating_now is True:
    q_input = st.chat_input("線上醫療咨詢服務。請輸入您的問題...")

#--------------------------------------------------------------------------------------------------------------
if "openai_model" not in st.session_state:
    st.session_state["openai_model"] = "gpt-3.5-turbo"

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"])

def set_moredesc():
    st.session_state.exec_moredesc = True      

#last_prompt = prompt
if prompt := q_input:
    st.session_state.last_prompt = "患者資料: 性別 {}, 年齡介於{}\n\n問題:".format(st.session_state.patient_six, st.session_state.patient_age) + prompt
    print("Prompt:",prompt)
    person_info = "患者資料: 性別 {}, 年齡介於{}\n\n".format(st.session_state.patient_six, st.session_state.patient_age)
    st.session_state.patient_info = person_info

    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        full_response = ""
        
        for response in client.chat.completions.create(
            model=st.session_state["openai_model"],
            max_tokens=256,
            temperature=0.1,
            top_p=1,
            frequency_penalty=0,
            messages=[
                {"role": m["role"], "content": "---------------------\n注意，簡單扼要針對患者的問題來回答，字數不多於128字，回答內容不可包含任何醫院名稱以及醫師姓名和其它與問題無關的內容。\n---------------------\n\n" + \
                    person_info + m["content"] }
                for m in st.session_state.messages
            ],
            stream=True,
        ):

            full_response += (response.choices[0].delta.content or "")
            message_placeholder.markdown(full_response + "▌")
            print(st.session_state.patient_info, end="     ")

        message_placeholder.markdown(full_response)
        st.session_state.last_response = full_response
        submit_moredesc = st.button(label='更詳細說明', on_click=set_moredesc)

        st.session_state.messages.append({"role": "assistant", "content": full_response})
        #--------------------------------------------------------------------------------------------------------------


if st.session_state.exec_moredesc:
    st.session_state.exec_moredesc = False
    more_ask = {"role": "user", "content": \
                "{}\n\n------------\n{}\n\n 你剛剛回答:{}\n------------\n病患希望你再針對剛簡短的回答內容，予以詳細的解釋說明，請不要重複之前的回答內容。相關醫療專有名詞請一起解釋。".format(st.session_state.patient_info, st.session_state.last_prompt, st.session_state.last_response) }
    #st.write(more_ask)
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        full_response = ""
        for response in client.chat.completions.create(
            model=st.session_state["openai_model"],
            max_tokens=2048,
            temperature=0.1,
            top_p=1,
            frequency_penalty=0,
            messages=[more_ask],
            stream=True,
        ):
            full_response += (response.choices[0].delta.content or "")
            message_placeholder.markdown(full_response + "▌")
            print(full_response, end="     ")

    st.session_state.messages.append({"role": "assistant", "content": full_response})