# -*- coding: utf-8 -*-

"""
@Time    :  2023/3/2 19:31
@Author  : 
@FileName: 
@Software: 
@Describe:
"""
#
# import redis
#
# redis_pool = redis.ConnectionPool(host='127.0.0.1', port=6379, password='', db=0)
# redis_conn = redis.Redis(connection_pool=redis_pool)
#
#
# name_dict = {
#     'name_4' : 'Zarten_4',
#     'name_5' : 'Zarten_5'
# }
# redis_conn.mset(name_dict)

import flask
import redis
import uuid
import json
from threading import Thread
import time

app = flask.Flask(__name__)
pool = redis.ConnectionPool(host='localhost', port=6379, max_connections=50)
redis_ = redis.Redis(connection_pool=pool, decode_responses=True)

db_key_query = 'query'
db_key_result = 'result'
batch_size = 32



def classify(batch_size):  # 调用模型，设置最大batch_size
    while True:
        texts = []
        query_ids = []
        if redis_.llen(db_key_query) == 0:  # 若队列中没有元素就继续获取
            continue
        for i in range(min(redis_.llen(db_key_query), batch_size)):
            query = redis_.lpop(db_key_query).decode('UTF-8')  # 获取query的text
            query_ids.append(json.loads(query)['id'])
            texts.append(json.loads(query)['text'])  # 拼接若干text 为batch
        result = texts  # 调用模型
        time.sleep(20)
        for (id_, res) in zip(query_ids, result):
            res['score'] = str(res['score'])
            redis_.set(id_, json.dumps(res))  # 将模型结果送回队列


@app.route("/predict", methods=["POST"])
def handle_query():
    text = flask.request.json['text']  # 获取用户query中的文本 例如"I love you"
    id_ = str(uuid.uuid1())  # 为query生成唯一标识
    d = {'id': id_, 'text': text}  # 绑定文本和query id
    redis_.rpush(db_key_query, json.dumps(d))  # 加入redis
    while True:
        result = redis_.get(id_)  # 获取该query的模型结果
        if result is not None:
            redis_.delete(id_)
            result_text = {'code': "200", 'data': result.decode('UTF-8')}
            break
    return flask.jsonify(result_text)  # 返回结果


if __name__ == "__main__":
    t = Thread(target=classify, args=(batch_size,))
    t.start()
    app.run(debug=False, host='127.0.0.1', port=9000)