import streamlit as st
import requests
from PIL import Image
from PIL.GifImagePlugin import GifImageFile
import os
# python -m streamlit run chat-ui.py

st.title("图书馆智能参考咨询AI BOT")
if 'generated' not in st.session_state:
    st.session_state['generated'] = []
if 'past' not in st.session_state:
    st.session_state['past'] = []
if 'ses' not in st.session_state:
    st.session_state['ses'] = requests.session()
if 'flag' not in st.session_state:
    st.session_state['flag'] = 0
ses = st.session_state['ses']

def generate_response(prompt):
    message = '暂时无法回答您的问题'
    if st.session_state['flag'] == 1:
        res = requests.post('http://localhost:44220/chat', json={'UserId': '18685121245', 'Content': prompt})
        data2 = res.json()
        if(data2['status'] == 1):
            st.session_state['flag'] = data2['flag']
            message = data2['output']
        else:
            st.session_state['flag'] = 0
    else:
        res = requests.post('http://localhost:8777/similarity', json={'text': prompt})
        data = res.json()
        print(data)
        if data['label'] in [0, 3]:
            message = data['answer']
        elif data['label'] in [1, 2, 4, 6, 7 ,8]:
            chat_content = ''
            if data['label'] == 4:
                chat_content = '报修%s' % prompt
            elif data['label'] == 6:
                chat_content = '建议%s' % prompt
            elif data['label'] == 7:
                chat_content = '投诉%s' % prompt
            else:
                chat_content = data['question']
            res = requests.post('http://localhost:44220/chat', json={'UserId': '18685121245', 'Content': chat_content})
            data2 = res.json()
            if(data2['status'] == 1):
                st.session_state['flag'] = data2['flag']
                message = data2['output']
            else:
                st.session_state['flag'] = 0
    return message

messages = st.container(height=400)
uploaded_file = st.file_uploader('选择文件', type=['pdf', 'docx'], accept_multiple_files=False)

if uploaded_file is not None:
    file_name = uploaded_file.name
    file_type = uploaded_file.type
    file_size = uploaded_file.size
    foler = os.path.dirname(__file__)
    filename = os.path.join(foler, 'data', file_name)
    with open(filename, "wb") as f:
        f.write(uploaded_file.getbuffer())
        f.close()
    print(filename)
    newfilename = filename.replace(':', '冒号').replace('\\', '斜杠').replace('-', '单杠').replace('.', '点点')
    output = generate_response(newfilename)
    st.session_state['past'].append(file_name)
    st.session_state['generated'].append(output)

if prompt := st.chat_input("Say something"):
    output = generate_response(prompt)
    st.session_state['past'].append(prompt)
    st.session_state['generated'].append(output)

if st.session_state['generated']:
    for i in range(len(st.session_state['generated'])):#-1, -1, -1
        messages.chat_message("你").write(st.session_state["past"][i])
        if type(st.session_state['generated'][i]) == str:
            answer = st.session_state['generated'][i]
            if answer.startswith('markdown:'):
                messages.markdown(answer[9:])
            else:
                messages.chat_message("图").write(st.session_state['generated'][i])
        elif type(st.session_state['generated'][i]) == list:
            for item in st.session_state['generated'][i]:
                messages.chat_message("图").write(item)