const express = require('express');
const axios = require('axios');

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

const TARGET_API = process.env.TARGET_API || "https://huggingface.co";

let PATH_MAPPINGS = { "/": "/" };
try {
    console.log(`Target API: ${TARGET_API}`);
    if (process.env.PATH_MAPPINGS) {
        PATH_MAPPINGS = JSON.parse(process.env.PATH_MAPPINGS);
    }
} catch (e) {
    console.error("解析 PATH_MAPPINGS 失败，使用默认配置:", e.message);
    PATH_MAPPINGS = { "/": "/" };
}

app.get('/', (req, res) => {
    res.send("service running.");
});

// 简单的工具：等待 ms 毫秒
const sleep = (ms) => new Promise(resolve => setTimeout(resolve, ms));

app.use(async (req, res) => {
    let fullPath = req.path;

    // 应用路径映射
    for (const [originalPath, newPath] of Object.entries(PATH_MAPPINGS)) {
        if (fullPath.startsWith(originalPath)) {
            fullPath = fullPath.replace(originalPath, newPath);
            break;
        }
    }

    const targetUrl = `${TARGET_API}${fullPath}`;

    // 复制并过滤掉 Host 请求头，防止目标服务器校验失败
    const headers = { ...req.headers };
    delete headers.host;

    // 为了支持 429 重试，需要把请求体缓存到内存中
    // （因为 req 流一旦被消费就无法再次使用）
    let bodyBuffer = null;
    const method = req.method.toUpperCase();
    if (['POST', 'PUT', 'PATCH', 'DELETE'].includes(method)) {
        const chunks = [];
        for await (const chunk of req) {
            chunks.push(chunk);
        }
        bodyBuffer = Buffer.concat(chunks);
    }

    // 监听客户端是否断开，避免无意义重试
    let aborted = false;
    const onAbort = () => { aborted = true; };
    res.on('close', onAbort);

    const MAX_RETRY = 20;
    let attempt = 0;

    try {
        while (attempt < MAX_RETRY) {
            attempt++;
            if (aborted) return;

            // 发起流式代理请求
            const response = await axios({
                method: req.method,
                url: targetUrl,
                headers: headers,
                params: req.query,
                data: bodyBuffer,        // 使用缓存的请求体，可重复发起
                responseType: 'stream',
                validateStatus: () => true
            });

            // 收到 429 且还有剩余重试次数时，销毁响应流并等待后重试
            if (response.status === 429 && attempt < MAX_RETRY) {
                // 释放当前响应流占用的资源
                try {
                    response.data.on('error', () => {});
                    response.data.destroy();
                } catch (e) {}

                const delay = 2000 + Math.random() * 3000; // 2 ~ 5 秒
                console.log(`收到 429，第 ${attempt} 次请求，将在 ${Math.round(delay)}ms 后重试（剩余 ${MAX_RETRY - attempt} 次）`);

                await sleep(delay);
                if (aborted) return;
                continue;
            }

            // 非 429 或已达最大重试次数：把响应透传给客户端
            res.status(response.status);

            const skipHeaders = ['content-length', 'transfer-encoding', 'connection'];
            Object.entries(response.headers).forEach(([key, value]) => {
                if (!skipHeaders.includes(key.toLowerCase())) {
                    res.setHeader(key, value);
                }
            });

            response.data.pipe(res);
            return;
        }
    } catch (error) {
        console.error(`代理请求失败: ${error.message}`);
        if (!res.headersSent) {
            res.status(500).send("Proxy Error");
        }
    } finally {
        res.off('close', onAbort);
    }
});

app.listen(PORT, '0.0.0.0', () => {
    console.log(`服务已在端口 ${PORT} 启动`);
});