import requests
import json
from flask import Flask, request, Response
import time
import configparser
import uuid
import threading
from cachetools import TTLCache

app = Flask(__name__)

# 读取配置文件
config = configparser.ConfigParser()
config.read('config.ini')

# 全局字典用于保存每个应用的 chatId，使用 api_key 作为键
# 使用 TTLCache 来设置每个 chatId 的有效期（例如：1小时）
# maxsize 设置缓存的最大容量，根据实际需求调整
chat_ids = TTLCache(maxsize=100, ttl=3600)  # 3600秒 = 1小时

# 创建一个锁以确保线程安全
chat_ids_lock = threading.Lock()


@app.route('/chat', methods=['POST'])
def chat():
    global chat_ids

    # 记录开始时间
    start_time = time.time()

    try:
        # 解析请求
        req_data = request.json
        dh_question = req_data.get('dh-question')
        dh_context = req_data.get('dh-context', "{}")
        uri = request.args.get('uri')
        api_key = request.args.get('bot_key')  # 从请求路径中获取 api_key

        if not api_key or not uri:
            return Response(json.dumps({"error": "Missing api_key or Missing uri"}), status=400,
                            mimetype='application/json')

        # 根据 api_key 生成或获取 chatId
        with chat_ids_lock:
            if api_key not in chat_ids or dh_context == "{}" or dh_context is None:
                # 如果是新对话，或者没有找到该 api_key 对应的 chatId，生成新的 chatId
                new_chat_id = str(uuid.uuid4())
                chat_ids[api_key] = new_chat_id  # 为该应用存储新的 chatId
            else:
                # 使用已有的 chatId
                new_chat_id = chat_ids[api_key]

        # 根据 api_key 设置 Authorization 头
        headers = {
            "Authorization": f"Bearer openapi-{api_key}",
            "Content-Type": "application/json"
        }

        # 构建请求数据
        data = {
            "chatId": new_chat_id,
            "stream": True,
            "detail": False,
            "messages": [
                {
                    "content": dh_question,
                    "role": "user"
                }
            ]
        }

        # 发送流式POST请求到fastGPT
        url = f"https://{uri}/api/v1/chat/completions"
        api_response = requests.post(url, headers=headers, data=json.dumps(data), stream=True)

        # 处理流式响应
        def generate():
            buffer = ""
            separators = "，。！？；"  # 分隔符列表
            if api_response.status_code == 200:
                for line in api_response.iter_lines():
                    if line:
                        decoded_line = line.decode('utf-8')
                        if decoded_line.startswith("data: "):
                            try:
                                json_data = json.loads(decoded_line[6:])
                                if json_data.get("choices"):
                                    delta = json_data["choices"][0].get("delta")
                                    if delta and delta.get("content"):
                                        content = delta["content"]
                                        buffer += content
                                        # 检查是否遇到分隔符
                                        for sep in separators:
                                            if sep in buffer:
                                                parts = buffer.split(sep)
                                                for part in parts[:-1]:
                                                    yield json.dumps({"answer": part + sep, "isEnd": False}) + "\n"
                                                buffer = parts[-1]
                                                break
                                    elif json_data["choices"][0].get("finish_reason") == "stop":
                                        if buffer:
                                            yield json.dumps({"answer": buffer, "isEnd": False}) + "\n"
                                        yield json.dumps({"answer": "", "isEnd": True}) + "\n"
                                        return  # 退出生成器函数
                            except json.JSONDecodeError as e:
                                if decoded_line[6:] != "[DONE]":
                                    print(f"JSONDecodeError: {e}, line: {decoded_line[6:]}")
                                else:
                                    if buffer:
                                        yield json.dumps({"answer": buffer, "isEnd": False}) + "\n"
                                    yield json.dumps({"answer": "", "isEnd": True}) + "\n"
                                    return  # 退出生成器函数
            else:
                yield json.dumps({"answer": f"Error: {api_response.status_code}", "isEnd": True}) + "\n"

        # 返回流式响应
        response = Response(generate(), content_type='application/json')

        # 记录结束时间并计算耗时
        end_time = time.time()
        elapsed_time = end_time - start_time
        print(f"Request processed in {elapsed_time:.2f} seconds")
        print(new_chat_id)

        return response

    except Exception as e:
        # 捕获所有异常并返回错误响应
        print(f"Exception occurred: {e}")
        return Response(json.dumps({"error": "Internal Server Error"}), status=500, mimetype='application/json')


if __name__ == '__main__':
    app.run(debug=True, port=5001, use_reloader=False)
------------
在流式处理部分增加容错逻辑，防止卡住