import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
from flask import Flask, jsonify
from flask import request
from transformers import pipeline
import redis
import uuid
import json
from threading import Thread
from vllm import LLM, SamplingParams
import time
import threading
import time
import concurrent.futures
import requests

app = Flask(__name__)
app.config["JSON_AS_ASCII"] = False
pool = redis.ConnectionPool(host='localhost', port=63179, max_connections=50,db=11, password="zhicheng123*")
redis_ = redis.Redis(connection_pool=pool, decode_responses=True)

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

sampling_params = SamplingParams(temperature=0.95, top_p=0.7,presence_penalty=0.9,stop="</s>", max_tokens=4096)
models_path = "/home/majiahui/project/models-llm/openbuddy-llama-7b-finetune"
llm = LLM(model=models_path, tokenizer_mode="slow")


def dialog_line_parse(url, text):
    """
    将数据输入模型进行分析并输出结果
    :param url: 模型url
    :param text: 进入模型的数据
    :return: 模型返回结果
    """

    response = requests.post(
        url,
        json=text,
        timeout=1000
    )
    if response.status_code == 200:
        return response.json()
    else:
        # logger.error(
        #     "【{}】 Failed to get a proper response from remote "
        #     "server. Status Code: {}. Response: {}"
        #     "".format(url, response.status_code, response.text)
        # )
        print("【{}】 Failed to get a proper response from remote "
            "server. Status Code: {}. Response: {}"
            "".format(url, response.status_code, response.text))
        print(text)
        return []


def classify(batch_size):  # 调用模型，设置最大batch_size
    while True:
        texts = []
        query_ids = []
        if redis_.llen(db_key_query) == 0:  # 若队列中没有元素就继续获取
            time.sleep(2)
            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
        outputs = llm.generate(texts, sampling_params)  # 调用模型

        generated_text_list = [""] * len(texts)
        for i, output in enumerate(outputs):
            index = output.request_id
            generated_text = output.outputs[0].text
            generated_text_list[int(index)] = generated_text


        for (id_, output) in zip(query_ids, generated_text_list):
            res = output
            redis_.set(id_, json.dumps(res))  # 将模型结果送回队列


@app.route("/predict", methods=["POST"])
def handle_query():
    text = request.json["texts"]  # 获取用户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': json.loads(result)}
            break
        time.sleep(1)

    return jsonify(result_text)  # 返回结果


@app.route("/articles_directory", methods=["POST"])
def articles_directory():
    text = request.json["texts"]  # 获取用户query中的文本 例如"I love you"
    nums = request.json["nums"]

    nums = int(nums)
    url = "http://114.116.25.228:18000/predict"

    input_data = []
    for i in range(nums):
        input_data.append([url, {"texts": text}])


    with concurrent.futures.ThreadPoolExecutor() as executor:
        # 使用submit方法将任务提交给线程池，并获取Future对象
        futures = [executor.submit(dialog_line_parse, i[0], i[1]) for i in input_data]

        # 使用as_completed获取已完成的任务，并获取返回值
        results = [future.result() for future in concurrent.futures.as_completed(futures)]

    return jsonify(results)  # 返回结果


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