const express = require("express");
const { io } = require("socket.io-client");
const { v4: uuidv4 } = require("uuid");
const { ProxyAgent } = require("proxy-agent");

// 创建代理实例
const 代理 = new ProxyAgent();

// 获取环境变量值并赋给中文变量名
const 曲奇池 = process.env.COOKIE_POOL || '';
const 用户代理 = process.env.USER_AGENT || '';
const 全局代理 = process.env.ALL_PROXY || '';

const 应用 = express();
const 端口 = process.env.PORT || 7860;

// 解析曲奇池，支持单个或多个曲奇（用逗号分隔）
const cookiePool = 曲奇池
    ? 曲奇池.split(',').map(cookie => cookie.trim()).filter(cookie => cookie.length > 0)
    : [];

if (cookiePool.length === 0) {
    console.error("错误：环境变量 曲奇池 未设置或为空。");
    process.exit(1);
}

// 创建事件流的工具函数
function 创建事件(事件, 数据) {
    if (typeof 数据 === "object") {
        数据 = JSON.stringify(数据);
    }
    return `event: ${事件}\ndata: ${数据}\n\n`;
}

// 处理 POST 请求的路由 /v1/messages
应用.post("/hf/v1/messages", (请求, 响应) => {
    请求.rawBody = "";
    请求.setEncoding("utf8");

    // 收集请求的原始数据
    请求.on("data", function (片段) {
        请求.rawBody += 片段;
    });

    请求.on("end", async () => {
        响应.setHeader("Content-Type", "text/event-stream;charset=utf-8");
        try {
            let 请求数据 = JSON.parse(请求.rawBody);
            if (请求数据.stream == false) {
                // 当 stream 为 false 时，返回固定的响应
                响应.send(
                    JSON.stringify({
                        id: uuidv4(),
                        content: [
                            {
                                text: "请开启流式传输。",
                            },
                            {
                                id: "string",
                                name: "string",
                                input: {},
                            },
                        ],
                        model: "string",
                        stop_reason: "end_turn",
                        stop_sequence: "string",
                        usage: {
                            input_tokens: 0,
                            output_tokens: 0,
                        },
                    })
                );
            } else if (请求数据.stream == true) {
                // 处理 stream 为 true 的情况
                let 用户消息 = [{ 问题: "", 回答: "" }];
                let 最后更新 = true;

                if (请求数据.system) {
                    // 将系统消息添加到消息列表的首位
                    请求数据.messages.unshift({ role: "system", content: 请求数据.system });
                }

                console.log(请求数据.messages);

                // 遍历消息列表，构建用户和助手的问答对
                请求数据.messages.forEach((消息) => {
                    if (消息.role == "system" || 消息.role == "user") {
                        if (最后更新) {
                            用户消息[用户消息.length - 1].问题 += 消息.content + "\n";
                        } else if (用户消息[用户消息.length - 1].问题 == "") {
                            用户消息[用户消息.length - 1].问题 += 消息.content + "\n";
                        } else {
                            用户消息.push({ 问题: 消息.content + "\n", 回答: "" });
                        }
                        最后更新 = true;
                    } else if (消息.role == "assistant") {
                        if (!最后更新) {
                            用户消息[用户消息.length - 1].回答 += 消息.content + "\n";
                        } else if (用户消息[用户消息.length - 1].回答 == "") {
                            用户消息[用户消息.length - 1].回答 += 消息.content + "\n";
                        } else {
                            用户消息.push({ 问题: "", 回答: 消息.content + "\n" });
                        }
                        最后更新 = false;
                    }
                });

                // 查找最新的用户消息
                const 最新用户消息 = 请求数据.messages
                    .slice().reverse()
                    .find(消息 => 消息.role === "user");

                if (!最新用户消息) {
                    throw new Error("请求中未找到用户消息。");
                }

                // 获取最新用户消息的索引
                const 最新用户索引 = 请求数据.messages.lastIndexOf(最新用户消息);

                // 构建除最新用户消息外的所有消息
                const 先前消息列表 = 请求数据.messages.slice(0, 最新用户索引);

                // 改进上下文拼接格式，增加提示并避免重复
                let 先前消息 = 先前消息列表
                    .map((消息) => {
                        if (消息.role === "user") {
                            return `用户: ${消息.content}`;
                        } else if (消息.role === "assistant") {
                            return `助手: ${消息.content}`;
                        } else if (消息.role === "system") {
                            return `系统: ${消息.content}`;
                        }
                        return 消息.content;
                    })
                    .join("\n\n");

                // 在上下文末尾添加提示，不再重复添加最新的用户消息
                先前消息 += `\n\n# ↓请关注用户最新的消息↓\n\n`;

                // 追加最新的用户消息
                先前消息 += `用户: ${最新用户消息.content}`;

                let 消息ID = uuidv4();

                // 发送消息开始的事件
                响应.write(
                    创建事件("message_start", {
                        type: "message_start",
                        message: {
                            id: 消息ID,
                            type: "message",
                            role: "assistant",
                            content: [],
                            model: "claude-3-opus-20240229",
                            stop_reason: null,
                            stop_sequence: null,
                            usage: { input_tokens: 8, output_tokens: 1 },
                        },
                    })
                );
                响应.write(创建事件("content_block_start", { type: "content_block_start", index: 0, content_block: { type: "text", text: "" } }));
                响应.write(创建事件("ping", { type: "ping" }));

                // 随机选择一个曲奇
                const 随机曲奇 = cookiePool[Math.floor(Math.random() * cookiePool.length)];

                // 配置 Socket.IO 客户端选项
                var 选项 = {
                    agent: 代理,
                    auth: {
                        jwt: "anonymous-ask-user",
                    },
                    reconnection: false,
                    transports: ["websocket"],
                    path: "/socket.io",
                    hostname: "www.perplexity.ai",
                    secure: true,
                    port: "443",
                    extraHeaders: {
                        Cookie: 随机曲奇, // 使用随机选择的曲奇
                        "User-Agent": 用户代理,
                        Accept: "*/*",
                        priority: "u=1, i",
                        Referer: "https://www.perplexity.ai/",
                    },
                };

                // 通过 Socket.IO 连接 Perplexity.ai 的 WebSocket
                var socket = io("wss://www.perplexity.ai/", 选项);

                socket.on("connect", function () {
                    console.log(" > [已连接]");
                    socket
                        .emitWithAck("perplexity_ask", 先前消息, {
                            "version": "2.9",
                            "source": "default",
                            "attachments": [],
                            "language": "en-GB",
                            "timezone": "Europe/London",
                            "search_focus": "writing",
                            "frontend_uuid": uuidv4(),
                            "mode": "concise",
                            "is_related_query": false,
                            "is_default_related_query": false,
                            "visitor_id": uuidv4(),
                            "frontend_context_uuid": uuidv4(),
                            "prompt_source": "user",
                            "query_source": "home"
                        })
                        .then((响应数据) => {
                            console.log(响应数据);
                            响应.write(创建事件("content_block_stop", { type: "content_block_stop", index: 0 }));
                            响应.write(
                                创建事件("message_delta", {
                                    type: "message_delta",
                                    delta: { stop_reason: "end_turn", stop_sequence: null },
                                    usage: { output_tokens: 12 },
                                })
                            );
                            响应.write(创建事件("message_stop", { type: "message_stop" }));

                            响应.end();
                        }).catch((错误) => {
                            if(错误.message != "socket has been disconnected"){
                                console.log(错误);
                            }
                        });
                });

                // 监听所有事件，记录日志
                socket.onAny((事件, ...参数) => {
                    console.log(`> [收到事件: ${事件}]`);
                });

                // 监听查询进度事件，处理返回的文本块
                socket.on("query_progress", (数据) => {
                    if(数据.text){
                        try {
                            var 文本 = JSON.parse(数据.text)
                            var 块 = 文本.chunks[文本.chunks.length - 1];
                            if (块) {
                                let 块JSON = JSON.stringify({
                                    type: "content_block_delta",
                                    index: 0,
                                    delta: { type: "text_delta", text: 块 },
                                });
                                响应.write(创建事件("content_block_delta", 块JSON));
                            }
                        } catch (错误) {
                            console.log("解析 query_progress 数据时出错:", 错误);
                        }
                    }
                });

                // 监听断开连接事件
                socket.on("disconnect", function () {
                    console.log(" > [已断开连接]");
                });

                // 监听错误事件，返回错误信息给客户端
                socket.on("error", (错误) => {
                    let 块JSON = JSON.stringify({
                        type: "content_block_delta",
                        index: 0,
                        delta: { type: "text_delta", text: "获取输出时出现错误\n请查看日志以获取更多信息" },
                    });
                    响应.write(创建事件("content_block_delta", 块JSON));
                    响应.write(创建事件("content_block_stop", { type: "content_block_stop", index: 0 }));
                    响应.write(
                        创建事件("message_delta", {
                            type: "message_delta",
                            delta: { stop_reason: "end_turn", stop_sequence: null },
                            usage: { output_tokens: 12 },
                        })
                    );
                    响应.write(创建事件("message_stop", { type: "message_stop" }));

                    响应.end();
                    console.log(错误);
                });

                // 监听连接错误事件，返回错误信息给客户端
                socket.on("connect_error", function (错误) {
                    let 块JSON = JSON.stringify({
                        type: "content_block_delta",
                        index: 0,
                        delta: { type: "text_delta", text: "连接到 Perplexity.ai 失败\n请查看日志以获取更多信息" },
                    });
                    响应.write(创建事件("content_block_delta", 块JSON));
                    响应.write(创建事件("content_block_stop", { type: "content_block_stop", index: 0 }));
                    响应.write(
                        创建事件("message_delta", {
                            type: "message_delta",
                            delta: { stop_reason: "end_turn", stop_sequence: null },
                            usage: { output_tokens: 12 },
                        })
                    );
                    响应.write(创建事件("message_stop", { type: "message_stop" }));

                    响应.end();
                    console.log(错误);
                });

                // 监听客户端关闭事件，断开 Socket 连接
                响应.on("close", function () {
                    console.log(" > [客户端已关闭连接]");
                    socket.disconnect();
                });
            } else {
                throw new Error("无效的请求");
            }
        } catch (错误) {
            console.log(错误);
            响应.write(JSON.stringify({ error: 错误.message }));
            响应.end();
            return;
        }
    });
});

// 处理所有未定义的路由，返回 418
应用.use((请求, 响应, 下一个) => {
    响应.status(418).send("418 I'm a teapot");
});

// 启动服务器，监听指定端口
应用.listen(端口, () => {
    console.log(`Perplexity 代理服务器正在监听端口 ${端口}`);
});