import requests
from flask import Flask, request, jsonify
import logging
import re

app = Flask(__name__)
logging.basicConfig(level=logging.DEBUG)


def preprocess_prompt(prompt):
    # 判断提示内容是否为纯英文
    if is_pure_english(prompt):
        processed_prompt = prompt
    else:
        # 使用正则表达式匹配 @startmindmap 和 @endmindmap 之间的内容
        pattern = r'@startmindmap\n(.*?)\n@endmindmap'
        # pattern = r'@startmindmap\n([*+\\-]*?)\n@endmindmap'
        match = re.search(pattern, prompt, re.DOTALL)
        if match:
            mindmap_content = match.group(1)
            # 将提示内容重复
            processed_prompt = f"@startmindmap\n{mindmap_content}\n{mindmap_content}\n@endmindmap"
        else:
            processed_prompt = prompt
    return processed_prompt


def is_pure_english(text):
    # 检查字符串中是否只包含英文字母、数字和一些特殊字符
    allowed_chars = set(
        'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789@\n*\-+ ')
    for char in text:
        if char not in allowed_chars:
            return False
    return True


@app.route('/v1/images/generations', methods=['POST'])
def generate_image():
    # 获取前端请求中的数据
    data = request.json
    logging.debug(f"Received request data: {data}")

    # 预处理 prompt
    processed_prompt = preprocess_prompt(data['prompt'])
    logging.debug(f"Processed prompt: {processed_prompt}")

    # 将数据发送到 https://plantumll.azurewebsites.net/coder 接口
    response = requests.post(
        'https://plantumll.azurewebsites.net/coder', data=processed_prompt)
    logging.debug(f"Response status code: {response.status_code}")
    logging.debug(f"Response text: {response.text}")

    # 处理响应数据
    response_data = response.text
    created = 1589478378  # 假设这是一个固定值
    url = f"https://plantumll.azurewebsites.net/png/{response_data}"

    # 构造返回给前端的数据
    result = {
        "created": created,
        "data": [
            {
                "url": url
            }
        ]
    }

    return jsonify(result)


if __name__ == '__main__':
    app.run(host='0.0.0.0',  port=7860)