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

const app = express();
const port = process.env.PORT || 8081;

// 从环境变量中获取 accessToken
const accessToken = process.env.PPLX_KEY;

console.log(`Server starting. Access token: ${accessToken ? "Set" : "Not set"}`);

// 添加一个中间件来解析 JSON 请求体
app.use(express.json());

// 添加一个中间件来记录所有请求
app.use((req, res, next) => {
    console.log(`[${new Date().toISOString()}] ${req.method} request received for ${req.path}`);
    console.log(`Headers: ${JSON.stringify(req.headers)}`);
    next();
});

var opts = {
    agent: agent,
    auth: {
        jwt: "anonymous-ask-user",
    },
    reconnection: false,
    transports: ["websocket"],
    path: "/socket.io",
    hostname: "www.perplexity.ai",
    secure: true,
    port: "443",
    extraHeaders: {
        Cookie: process.env.PPLX_COOKIE,
        "User-Agent": process.env.USER_AGENT,
        Accept: "*/*",
        priority: "u=1, i",
        Referer: "https://www.perplexity.ai/",
    },
};

app.post("/v1/messages", (req, res) => {
    console.log("[POST /v1/messages] Processing request");
    
    // 严格验证客户端的 Authorization 头
    const clientAuth = req.headers['authorization'];
    console.log(`[POST /v1/messages] Received Authorization header: ${clientAuth}`);
    console.log(`[POST /v1/messages] Expected Authorization: Bearer ${accessToken}`);
    
    if (!clientAuth || clientAuth !== `Bearer ${accessToken}`) {
        console.log("[POST /v1/messages] Authorization failed. Sending 401 Unauthorized response.");
        res.status(401).json({ error: "Unauthorized" });
        return;
    }
    
    console.log("[POST /v1/messages] Authorization successful. Proceeding with request processing.");

    // 使用解析后的 JSON 体
    const jsonBody = req.body;
    console.log(`[POST /v1/messages] Received request body: ${JSON.stringify(jsonBody)}`);

    res.setHeader("Content-Type", "text/event-stream;charset=utf-8");
    try {
        if (jsonBody.stream == false) {
            console.log("[POST /v1/messages] Stream is false. Sending non-streaming response.");
            res.send(
                JSON.stringify({
                    id: uuidv4(),
                    content: [
                        {
                            text: "Please turn on streaming.",
                        },
                        {
                            id: "string",
                            name: "string",
                            input: {},
                        },
                    ],
                    model: "string",
                    stop_reason: "end_turn",
                    stop_sequence: "string",
                    usage: {
                        input_tokens: 0,
                        output_tokens: 0,
                    },
                })
            );
        } else if (jsonBody.stream == true) {
            console.log("[POST /v1/messages] Stream is true. Processing streaming request.");
            // 计算用户消息长度
            let userMessage = [{ question: "", answer: "" }];
            let lastUpdate = true;
            if (jsonBody.system) {
                // 把系统消息加入messages的首条
                jsonBody.messages.unshift({ role: "system", content: jsonBody.system });
            }
            console.log(`[POST /v1/messages] Processed messages: ${JSON.stringify(jsonBody.messages)}`);
            jsonBody.messages.forEach((msg) => {
                if (msg.role == "system" || msg.role == "user") {
                    if (lastUpdate) {
                        userMessage[userMessage.length - 1].question += msg.content + "\n";
                    } else if (userMessage[userMessage.length - 1].question == "") {
                        userMessage[userMessage.length - 1].question += msg.content + "\n";
                    } else {
                        userMessage.push({ question: msg.content + "\n", answer: "" });
                    }
                    lastUpdate = true;
                } else if (msg.role == "assistant") {
                    if (!lastUpdate) {
                        userMessage[userMessage.length - 1].answer += msg.content + "\n";
                    } else if (userMessage[userMessage.length - 1].answer == "") {
                        userMessage[userMessage.length - 1].answer += msg.content + "\n";
                    } else {
                        userMessage.push({ question: "", answer: msg.content + "\n" });
                    }
                    lastUpdate = false;
                }
            });
            // user message to plaintext
            let previousMessages = jsonBody.messages
                .map((msg) => {
                    return msg.content
                })
                .join("\n\n");

            console.log(`[POST /v1/messages] Previous messages: ${previousMessages}`);

            let msgid = uuidv4();
            // send message start
            res.write(
                createEvent("message_start", {
                    type: "message_start",
                    message: {
                        id: msgid,
                        type: "message",
                        role: "assistant",
                        content: [],
                        model: "claude-3-opus-20240229",
                        stop_reason: null,
                        stop_sequence: null,
                        usage: { input_tokens: 8, output_tokens: 1 },
                    },
                })
            );
            res.write(createEvent("content_block_start", { type: "content_block_start", index: 0, content_block: { type: "text", text: "" } }));
            res.write(createEvent("ping", { type: "ping" }));

            console.log("[POST /v1/messages] Initiating WebSocket connection to Perplexity.ai");
            // proxy response
            var socket = io("wss://www.perplexity.ai/", opts);

            socket.on("connect", function () {
                console.log("[POST /v1/messages] WebSocket connected to Perplexity.ai");
                socket
                    .emitWithAck("perplexity_ask", previousMessages, {
                        "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((response) => {
                        console.log(`[POST /v1/messages] Received response from Perplexity.ai: ${JSON.stringify(response)}`);
                        res.write(createEvent("content_block_stop", { type: "content_block_stop", index: 0 }));
                        res.write(
                            createEvent("message_delta", {
                                type: "message_delta",
                                delta: { stop_reason: "end_turn", stop_sequence: null },
                                usage: { output_tokens: 12 },
                            })
                        );
                        res.write(createEvent("message_stop", { type: "message_stop" }));

                        res.end();
                    }).catch((error) => {
                        if(error.message != "socket has been disconnected"){
                            console.log(`[POST /v1/messages] Error in WebSocket communication: ${error}`);
                        }
                    });
            });
            socket.onAny((event, ...args) => {
                console.log(`[POST /v1/messages] Received WebSocket event: ${event}`);
            });
            socket.on("query_progress", (data) => {
                if(data.text){
                    var text = JSON.parse(data.text)
                    var chunk = text.chunks[text.chunks.length - 1];
                    if(chunk){
                        console.log(`[POST /v1/messages] Received chunk: ${chunk}`);
                        chunkJSON = JSON.stringify({
                            type: "content_block_delta",
                            index: 0,
                            delta: { type: "text_delta", text: chunk },
                        });
                        res.write(createEvent("content_block_delta", chunkJSON));
                    }
                }
            });
            socket.on("disconnect", function () {
                console.log("[POST /v1/messages] WebSocket disconnected from Perplexity.ai");
            });
            socket.on("error", (error) => {
                console.log(`[POST /v1/messages] WebSocket error: ${error}`);
                chunkJSON = JSON.stringify({
                    type: "content_block_delta",
                    index: 0,
                    delta: { type: "text_delta", text: "Error occured while fetching output 输出时出现错误\nPlease refer to the log for more information 请查看日志以获取更多信息" },
                });
                res.write(createEvent("content_block_delta", chunkJSON));
                res.write(createEvent("content_block_stop", { type: "content_block_stop", index: 0 }));
                res.write(
                    createEvent("message_delta", {
                        type: "message_delta",
                        delta: { stop_reason: "end_turn", stop_sequence: null },
                        usage: { output_tokens: 12 },
                    })
                );
                res.write(createEvent("message_stop", { type: "message_stop" }));

                res.end();
            });
            socket.on("connect_error", function (error) {
                console.log(`[POST /v1/messages] WebSocket connection error: ${error}`);
                chunkJSON = JSON.stringify({
                    type: "content_block_delta",
                    index: 0,
                    delta: { type: "text_delta", text: "Failed to connect to the Perplexity.ai 连接到Perplexity失败\nPlease refer to the log for more information 请查看日志以获取更多信息" },
                });
                res.write(createEvent("content_block_delta", chunkJSON));
                res.write(createEvent("content_block_stop", { type: "content_block_stop", index: 0 }));
                res.write(
                    createEvent("message_delta", {
                        type: "message_delta",
                        delta: { stop_reason: "end_turn", stop_sequence: null },
                        usage: { output_tokens: 12 },
                    })
                );
                res.write(createEvent("message_stop", { type: "message_stop" }));

                res.end();
            });
            res.on("close", function () {
                console.log("[POST /v1/messages] Client closed connection");
                socket.disconnect();
            });
        } else {
            console.log("[POST /v1/messages] Invalid request: stream is neither true nor false");
            throw new Error("Invalid request");
        }
    } catch (e) {
        console.log(`[POST /v1/messages] Error in request processing: ${e}`);
        res.write(JSON.stringify({ error: e.message }));
        res.end();
        return;
    }
});

// 处理 /ai/v1/messages 路由
app.post("/ai/v1/messages", (req, res) => {
    console.log("[POST /ai/v1/messages] Received request, forwarding to /v1/messages");
    app.handle(req, res, req.next);
});

// handle other
app.use((req, res, next) => {
    console.log(`[${new Date().toISOString()}] Received request for ${req.path} - returning 404`);
    res.status(404).send("Not Found");
});

app.listen(port, () => {
    console.log(`[${new Date().toISOString()}] Perplexity proxy listening on port ${port}`);
});

// eventStream util
function createEvent(event, data) {
    // if data is object, stringify it
    if (typeof data === "object") {
        data = JSON.stringify(data);
    }
    return `event: ${event}\ndata: ${data}\n\n`;
}