from time import sleep
import erniebot

from flask import Flask, request, jsonify, stream_with_context, send_file
import cv2
import os
from PIL import Image
from paddleocr import PaddleOCR, draw_ocr
import numpy as np
from gevent import pywsgi
from flask_cors import CORS
import hashlib

app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 设置请求体最大允许 16MB
CORS(app)
model = PaddleOCR(use_angle_cls=True, debug=False, ocr_version="PP-OCRv4")


def md5sum(file):
    m = hashlib.md5()
    m.update(file)
    ret = m.hexdigest()
    return ret


@app.route('/upload', methods=['get', 'post'])
def upload():
    username = request.form.get("username")
    img = request.files['file']
    picname = img.filename
    suffix = picname.split(".")[-1]
    file = img.read()
    md5 = md5sum(file)  # 使用md5作为文件名避免文件重复上传

    # 创建用户文件夹
    imgfile1_path = "./img/" + username + "/"
    if not os.path.exists(imgfile1_path):
        os.makedirs(imgfile1_path)

    img1_path = os.path.join(imgfile1_path, md5) + "." + str(suffix)
    if not os.path.exists(img1_path):
        file = cv2.imdecode(np.frombuffer(file, np.uint8), cv2.IMREAD_COLOR)  # 解码为ndarray
        cv2.imwrite(filename=img1_path, img=file)

    tempmap = {"path": img1_path}
    return jsonify(tempmap)


@app.route('/img/<username>/<pic_name>', methods=["GET"])
def get_img(username, pic_name):
    img_path = os.path.join("./img", username, pic_name)
    return send_file(img_path)


@app.route('/img/<pic_name>', methods=['GET'])
def get_image(pic_name):
    img_path = os.path.join("./img", pic_name)
    return send_file(img_path)


@app.route("/ocr", methods=["POST"])
def ocr():
    img_path = request.form.get("path")
    img_arr = np.array(Image.open(img_path))
    results = model.ocr(img_arr)
    texts = []
    for result in results[0]:
        texts.append(result[1][0])
    response = {
        "texts": texts
    }
    return response


def set_token(token):
    erniebot.api_type = 'aistudio'
    erniebot.access_token = token


def generate(response):
    for resp in response:
        yield resp.get_result()


def get_ernie_response(prompt):
    # model_name = 'ernie-3.5'
    model_name = 'ernie-4.0'
    messages = [{'role': 'user', 'content': prompt}]
    response = erniebot.ChatCompletion.create(
        model=model_name,
        messages=messages,
        stream=True,
        system="你不需要回答其他信息, 仅返回结果即可"
    )
    return response


@app.route("/polish", methods=["POST"])
def polish():
    prompt = ("你需要根据所给出笔记的场景对笔记进行润色"
              "给出的笔记是：{}\n").format(request.form.get("msg"))
    response = get_ernie_response(prompt)
    return app.response_class(stream_with_context((generate(response))))


@app.route("/abstract", methods=["POST"])
def abstract():
    prompt = ("接下来我会给出我的笔记,你需要提炼,缩写我的笔记,尽量精简,要比之前的笔记字数少,"
              "你需要将最重要的那些信息使用html标签进行加粗、斜体、加红色等; 最终不需要回答其他信息，返回结果即可."
              "给出的笔记是：{}\n").format(request.form.get("msg"))
    response = get_ernie_response(prompt)
    return app.response_class(stream_with_context(generate(response)))


@app.route("/expand", methods=["POST"])
def expand():
    prompt = ("接下来我会给出我的笔记,你需要根据关键信息扩写我的笔记,尽量详细."
              "要求: 1.你需要将最重要的那些信息使用html标签进行加粗、斜体、加红色等;"
              "2.充分使用多级标题,使用<h1>, <h2>, <h3>等标签;"
              "3.最终不需要回答其他信息,返回结果即可.给出的笔记是: {}").format(request.form.get("msg"))
    response = get_ernie_response(prompt)
    return app.response_class(stream_with_context(generate(response)))


@app.route("/segment", methods=["POST"])
def segment():
    prompt = ("接下来我会给出一篇笔记, 请将其根据内容使用各级标题进行合理的分段.\n "
              "要求：1. 你需要将最重要的那些信息使用html标签进行加粗、斜体、加红色等;"
              "2.充分使用多级标题, 使用<h1>, <h2>, <h3>等标签;"
              "3.最终不需要回答其他信息，返回结果即可。给出的笔记是：{}\n").format(request.form.get("msg"))
    response = get_ernie_response(prompt)
    return app.response_class(stream_with_context(generate(response)))


@app.route("/translate", methods=["POST"])
def translate():
    prompt = "对下面这段话进行翻译,如果是中文则翻译为英文,如果是英文则翻译为中文,仅回答结果即可不需要其他信息: {}".format(
        request.form.get("msg"))
    response = get_ernie_response(prompt)
    return app.response_class(stream_with_context(generate(response)))


@app.route("/table", methods=["POST"])
def generateTable():
    prompt = ("请根据我的笔记，生成结构清晰的html格式的表格."
              "要求：1、保持表格的清晰排列和标签"
              "2、根据笔记的内容主旨生成简短扼要的表格标题"
              "3、表格内容能够保持和表格列标题一致时，用一个表格概括；否则可以用多个表格概括，但必须都要保持表格的格式"
              "4、使用html的语法创建表格"
              "5、最终不需要回答其他信息, 也不要返回```html代码块, 仅返回<table>部分的html代码即可" +
              "给出的笔记是: {}").format(request.form.get("msg"))
    response = get_ernie_response(prompt)
    return app.response_class(stream_with_context(generate(response)))


@app.route("/format", methods=["POST"])
def formatText():
    prompt = ("根据以下笔记的内容对其使用html标签进行排版,"
              "可以使用<b>、<p>、<i>、<u>、<mark>、<span>、<ul>、<li>、<h1>到<h4>等标签,"
              "可以使用text-align、color等style,仅回答结果即可,不要给出其他内容,给出的笔记是:{}").format(request.form.get("msg"))
    response = get_ernie_response(prompt)
    return app.response_class(stream_with_context(generate(response)))


if __name__ == "__main__":
    set_token("7587cca675d824db4ec5d7fe22c07bf5dac59a45")
    # server = pywsgi.WSGIServer(("127.0.0.1", 5000), app)
    server = pywsgi.WSGIServer(("0.0.0.0", 5000), app)
    server.serve_forever()