#!/usr/bin/env python
# encoding: utf-8
import gradio as gr
from PIL import Image
import requests
import base64
from io import BytesIO
import traceback
import os


def upload_img(image,_chatbot,_app_session):
    image = Image.fromarray(image)
    _app_session['sts']=None
    _app_session['ctx']=''
    _app_session['img']=image
    _chatbot.append(('图片解析成功，可以和我对话了', ''))
    return _chatbot,_app_session


def respond( _question, _chat_bot,_app_cfg):
    try:
        img = _app_cfg['img']
        buffered = BytesIO()
        img.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
        url = os.environ['SERVICE_URL']
        resp = requests.post(url, headers={
            "X-Model-Best-Model": "viscpm-chat-balance",
            "X-Model-Best-Trace-ID": "test-trace",
        }, json={
            "image": img_str,
            "question": _question,
        })
        resp = resp.json()
        # _answer = resp['data']['response']
        print('get response', resp)
        _answer = resp['response']
        print(f'question: {_question}, answer: {_answer}')
    except Exception as e:
        print(traceback.format_exc())
        if resp is not None:
            print(resp.content)
        _answer = "请求失败"
    _chat_bot.append((_question, _answer))
    _context = _app_cfg['ctx'] + '\n' + _question + '\n' + _answer + '\n'
    sts = None
    _app_cfg['ctx'] = _context
    _app_cfg['sts'] = sts
    return '', _chat_bot, _app_cfg


with gr.Blocks() as demo:
    app_session = gr.State({'sts':None,'ctx':None,'img':None})
    bt_pic = gr.Image(label="先上传一张图片")
    chat_bot = gr.Chatbot(label="聊天对话")
    txt_message = gr.Textbox(label="输入文字")

    txt_message.submit(respond, [ txt_message, chat_bot,app_session], [txt_message,chat_bot,app_session])
    bt_pic.upload(lambda: None, None, chat_bot, queue=False).then(upload_img, inputs=[bt_pic,chat_bot,app_session], outputs=[chat_bot,app_session])


demo.queue(concurrency_count=1, max_size=20).launch(share=False, debug=True)