import json
import uuid  # 用于生成随机用户 ID

import requests
from flask import Flask, request, Response

app = Flask(__name__)
# 全局字典用于保存用户的上下文对话
user_contexts = {}
MAX_HISTORY_LENGTH = 15  # 最大上下文历史长度


def get_models():
    models = {
        "object": "list",
        "data": [
            {"id": "Qwen2.5-72B", "object": "model", "created": 0, "owned_by": "Qwen"},
            {"id": "Llama3.3-70B", "object": "model", "created": 0, "owned_by": "Nemotron"},
            {"id": "Pixtral-124B", "object": "model", "created": 0, "owned_by": "Pixtral"},
            {"id": "Qwen2.5-Coder-32B", "object": "model", "created": 0, "owned_by": "Qwen"},
        ]
    }
    return json.dumps(models)


# 新的方法，检查model是否有效并返回id
def get_model_by_id(model_id=None):
    # 获取所有模型数据
    models_data = json.loads(get_models())["data"]
    
    # 提取所有有效的模型ID
    valid_ids = [model["id"] for model in models_data]
    
    # 如果model_id无效，默认返回Qwen2.5-72B
    if model_id not in valid_ids:
        model_id = "Qwen2.5-72B"
    
    # 根据model_id获取对应的模型数据
    model_data = next((model for model in models_data if model["id"] == model_id), None)
    
    # 返回model_data的id字段
    return model_data["id"] if model_data else None

def chat_completion(
        user_prompt, user_id: str = None, system_prompt="You are a helpful assistant.", model="Qwen2.5-72B",
        project="DecentralGPT", stream=False, temperature=0.3, max_tokens=1024, top_p=0.5,
        frequency_penalty=0, presence_penalty=0):
    """处理用户请求并保留上下文"""
    url = 'https://usa-chat.degpt.ai/api/v0/chat/completion/proxy'
    headers = {
        'accept': 'application/json',
        'accept-language': 'en-US,en;q=0.9,zh-CN;q=0.8,zh;q=0.7',
        'content-type': 'application/json',
        'dnt': '1',
        'origin': 'https://www.degpt.ai',
        'priority': 'u=1, i',
        'referer': 'https://www.degpt.ai/',
        'sec-ch-ua': 'Google Chrome";v="131", "Chromium";v="131", "Not_A Brand";v="24"',
        'sec-ch-ua-mobile': '?0',
        'sec-ch-ua-platform': '"macOS"',
        'sec-fetch-dest': 'empty',
        'sec-fetch-mode': 'cors',
        'sec-fetch-site': 'same-site',
        'user-agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36'
    }

    # 初始化或更新用户的对话历史
    if user_id is not None:
        if user_id not in user_contexts:
            user_contexts[user_id] = [{"role": "system", "content": system_prompt}]

        user_contexts[user_id].append({"role": "user", "content": user_prompt})

        # 检查是否需要修剪历史记录，保留 `system` 提示词
        while len(user_contexts[user_id]) > MAX_HISTORY_LENGTH:
            # 删除最早的用户问题和系统回复，但保留 `system` 提示词
            if len(user_contexts[user_id]) > 2:
                # 检查删除的条目是否有匹配的系统回复，如果没有，只删除用户输入
                if user_contexts[user_id][2]["role"] == "user":
                    user_contexts[user_id] = [user_contexts[user_id][0]] + user_contexts[user_id][2:]
                else:
                    user_contexts[user_id] = [user_contexts[user_id][0]] + user_contexts[user_id][2:]
            else:
                break

        messages = user_contexts[user_id]
    else:
        # 如果没有提供 user_id，不保留上下文
        messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}]

    payload = {
        "model": model,
        "messages": messages,
        "project": project,
        "stream": stream,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "top_p": top_p,
        "frequency_penalty": frequency_penalty,
        "presence_penalty": presence_penalty
    }

    try:
        response = requests.post(url, headers=headers, json=payload)
        response.encoding = 'utf-8'
        response.raise_for_status()
        ## print(response.text)

        # 获取响应并添加到上下文
        response_content = response.json()["choices"][0]["message"]["content"]
        # print(
        #     f"=========== {user_id}:{user_prompt} ====================\r\n请求内容:{messages}\r\n完整响应:{response.text}")

        # 将系统的回复添加到用户上下文中
        if user_id is not None:
            user_contexts[user_id].append({"role": "assistant", "content": response_content})
        return response.text
    except requests.exceptions.RequestException as e:
        print(f"请求失败: {e}")
        return "请求失败，请检查网络或参数配置。"
    except (KeyError, IndexError) as e:
        print(f"解析响应时出错: {e}")
        return "解析响应内容失败。"
    return {}

def chat_completion_messages(
        messages, user_id: str = None, system_prompt="You are a helpful assistant.", model="Qwen2.5-72B",
        project="DecentralGPT", stream=False, temperature=0.3, max_tokens=1024, top_p=0.5,
        frequency_penalty=0, presence_penalty=0):
    """处理用户请求并保留上下文"""
    url = 'https://usa-chat.degpt.ai/api/v0/chat/completion/proxy'
    headers = {
        'accept': 'application/json',
        'accept-language': 'en-US,en;q=0.9,zh-CN;q=0.8,zh;q=0.7',
        'content-type': 'application/json',
        'dnt': '1',
        'origin': 'https://www.degpt.ai',
        'priority': 'u=1, i',
        'referer': 'https://www.degpt.ai/',
        'sec-ch-ua': 'Google Chrome";v="131", "Chromium";v="131", "Not_A Brand";v="24"',
        'sec-ch-ua-mobile': '?0',
        'sec-ch-ua-platform': '"macOS"',
        'sec-fetch-dest': 'empty',
        'sec-fetch-mode': 'cors',
        'sec-fetch-site': 'same-site',
        'user-agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36'
    }

    # # 初始化或更新用户的对话历史
    # if user_id is not None:
    #     if user_id not in user_contexts:
    #         user_contexts[user_id] = [{"role": "system", "content": system_prompt}]

    #     user_contexts[user_id].append({"role": "user", "content": user_prompt})

    #     # 检查是否需要修剪历史记录，保留 `system` 提示词
    #     while len(user_contexts[user_id]) > MAX_HISTORY_LENGTH:
    #         # 删除最早的用户问题和系统回复，但保留 `system` 提示词
    #         if len(user_contexts[user_id]) > 2:
    #             # 检查删除的条目是否有匹配的系统回复，如果没有，只删除用户输入
    #             if user_contexts[user_id][2]["role"] == "user":
    #                 user_contexts[user_id] = [user_contexts[user_id][0]] + user_contexts[user_id][2:]
    #             else:
    #                 user_contexts[user_id] = [user_contexts[user_id][0]] + user_contexts[user_id][2:]
    #         else:
    #             break

    #     messages = user_contexts[user_id]
    # else:
    #     # 如果没有提供 user_id，不保留上下文
    #     messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}]

    payload = {
        "model": model,
        "messages": messages,
        "project": project,
        "stream": stream,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "top_p": top_p,
        "frequency_penalty": frequency_penalty,
        "presence_penalty": presence_penalty
    }

    try:
        response = requests.post(url, headers=headers, json=payload)
        response.encoding = 'utf-8'
        response.raise_for_status()
        ## print(response.text)

        # # 获取响应并添加到上下文
        # response_content = response.json()["choices"][0]["message"]["content"]
        # # print(
        # #     f"=========== {user_id}:{user_prompt} ====================\r\n请求内容:{messages}\r\n完整响应:{response.text}")

        # # 将系统的回复添加到用户上下文中
        # if user_id is not None:
        #     user_contexts[user_id].append({"role": "assistant", "content": response_content})
        return response.text
    except requests.exceptions.RequestException as e:
        print(f"请求失败: {e}")
        return "请求失败，请检查网络或参数配置。"
    except (KeyError, IndexError) as e:
        print(f"解析响应时出错: {e}")
        return "解析响应内容失败。"
    return {}
@app.route('/api/models', methods=['GET'])
@app.route('/api/v1/models', methods=['GET'])
@app.route('/hf/v1/models', methods=['GET'])
def models():
    """返回可用模型列表"""
    return get_models()


@app.route('/api/chat/completion', methods=['POST'])
@app.route('/api/v1/chat/completions', methods=['POST'])
@app.route('/hf/v1/chat/completions', methods=['POST'])
def chat_completion_api():
    """处理用户请求并保留上下文"""
    data = request.json
    # user_prompt = data.get("prompt")
    user_id = data.get("user_id", str(uuid.uuid4()))  # 如果未提供 user_id，生成随机值
    model=get_model_by_id(data.get("model"))
    temperature=data.get("temperature", 0.3)
    max_tokens =data.get("max_tokens", 1024)
    messages = data.get("messages", [])
    stream = data.get("stream", False)
    response_content = chat_completion_messages(
        messages=messages,
        user_id=user_id,
        temperature=temperature,
        max_tokens=max_tokens,
        stream=stream,
        model=model
    )

    # maybe \uxxxx
    # return jsonify(response_content)
    ## maybe \"xxx\"
    # return Response(
    #     json.dumps(response_content, ensure_ascii=False),
    #     content_type="application/json; charset=utf-8"
    # )
    # support Chinese
    if isinstance(response_content, str):  # 如果已经是 JSON 字符串
        return Response(response_content, content_type="application/json; charset=utf-8")


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