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

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

// 从环境变量获取目标 API 基础 URL
const TARGET_API = process.env.TARGET_API || "https://huggingface.co";

// 1. 客户端访问此代理所需的 API Key（若未设置则不进行校验）
const AUTH_API_KEY = process.env.AUTH_API_KEY;

// 2. 解析请求目标服务器时需要注入的 API Key 数组
let DEST_API_KEYS = [];
if (process.env.DEST_API_KEY) {
    try {
        DEST_API_KEYS = JSON.parse(process.env.DEST_API_KEY);
        if (!Array.isArray(DEST_API_KEYS)) {
            DEST_API_KEYS = [process.env.DEST_API_KEY];
        }
    } catch (e) {
        DEST_API_KEYS = process.env.DEST_API_KEY.split(',').map(k => k.trim()).filter(Boolean);
    }
}

// 解析路径映射环境变量
let PATH_MAPPINGS = { "/": "/" };
try {
    if (process.env.PATH_MAPPINGS) {
        PATH_MAPPINGS = JSON.parse(process.env.PATH_MAPPINGS);
    }
} catch (e) {
    console.error("解析 PATH_MAPPINGS 失败，使用默认配置:", e.message);
    PATH_MAPPINGS = { "/": "/" };
}

// 明确定义的路由放在最上面：根路径 GET 请求直接响应（用于健康检查，免 Key 校验）
app.get('/', (req, res) => {
    res.send("service running.");
});

// 使用不带路径参数的 app.use() 作为万能兜底代理
app.use(async (req, res) => {
    
    // 校验客户端的 AUTH_API_KEY
    if (AUTH_API_KEY) {
        let clientKey = '';

        // 判断路径是否包含 /v1/messages 或 /v1/responses
        if (req.path.includes('/v1/messages') || req.path.includes('/v1/responses')) {
            // 优先从 x-api-key 获取，没有则尝试从 Authorization 获取
            clientKey = req.headers['x-api-key'] || (req.headers['authorization'] || '').replace(/^Bearer\s+/i, '').trim();
        } else {
            // 其他路径走标准的 Authorization Bearer 认证
            const authHeader = req.headers['authorization'] || '';
            clientKey = authHeader.replace(/^Bearer\s+/i, '').trim();
        }

        if (clientKey !== AUTH_API_KEY) {
            return res.status(401).send("Unauthorized: Invalid or missing API Key");
        }
    }

    let fullPath = req.path;
    // 应用路径映射（仅替换一次）
    for (const [originalPath, newPath] of Object.entries(PATH_MAPPINGS)) {
        if (fullPath.startsWith(originalPath)) {
            fullPath = fullPath.replace(originalPath, newPath);
            break;
        }
    }

    // 拼接最终目标 URL
    const targetUrl = `${TARGET_API}${fullPath}`;

    // 复制请求头并过滤掉 host
    const headers = { ...req.headers };
    delete headers.host;
    
    // 请求体缓存：为了支持重试，必须将 req 流缓存为 Buffer，否则流只能读一次，重试时将没有数据可发
    let bodyBuffer = [];
    req.on('data', chunk => bodyBuffer.push(chunk));
    // 等待请求体接收完毕再执行后续逻辑
    await new Promise(resolve => req.on('end', resolve));
    const bodyData = Buffer.concat(bodyBuffer);

    // 根据目标路径动态注入认证头
    if (DEST_API_KEYS.length > 0) {
        const randomKey = DEST_API_KEYS[Math.floor(Math.random() * DEST_API_KEYS.length)];
        
        // 检查映射后的路径并细分处理
        if (fullPath.includes('messages')) {
            // messages 路径：仅使用 x-api-key
            headers['x-api-key'] = randomKey;
            delete headers['authorization']; // 强行移除原 Authorization，防止双认证头冲突
        } else if (fullPath.includes('responses')) {
            // responses 路径：同时保留并设置两个认证头
            headers['x-api-key'] = randomKey;
            headers['authorization'] = `Bearer ${randomKey}`;
        } else {
            // 传统 /v1/chat 等路由，走 OpenAI 标准的 Authorization 头
            headers['authorization'] = `Bearer ${randomKey}`;
            delete headers['x-api-key'];
        }
    }

    // 【排查日志】打印每次转发的真实细节
    //console.log(`[Proxy Request] Method: ${req.method} | URL: ${targetUrl}`);
    //console.log(`[Proxy Headers] Authorization: ${headers['authorization'] ? 'Present' : 'None'} | x-api-key: ${headers['x-api-key'] ? 'Present' : 'None'}`);

    const MAX_RETRIES = 60;
    let attempt = 0;

    // 重试循环
    while (attempt < MAX_RETRIES) {
        attempt++;
        try {
            // 发起代理请求
            const response = await axios({
                method: req.method,
                url: targetUrl,
                headers: headers,
                params: req.query,
                data: bodyData.length > 0 ? bodyData : undefined, // 使用缓存的请求体
                responseType: 'stream', 
                validateStatus: () => true // 允许捕获所有状态码，交给我们手动判断
            });

            // 打印目标服务器返回的状态码
            //console.log(`[Proxy Response] Attempt ${attempt} | Status Code From Target: ${response.status}`);

            // 如果状态码是 200，则将其响应流管道输送给客户端并退出循环
            if (response.status === 200) {
                // 复制响应状态码
                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; // 成功响应，退出函数
            } else {
                // 状态码不是 200，销毁当前流，进入下一次重试
                response.data.destroy(); 
                
                // 如果已达最大重试次数，直接把最后一次的非200状态码透传给客户端并返回错误信息
                if (attempt >= MAX_RETRIES) {
                    console.error(`[Proxy Failed] Reached max retries (${MAX_RETRIES}). Last status: ${response.status}`);
                    if (!res.headersSent) {
                        res.status(response.status).send(`Proxy Error: Target server returned status ${response.status} after ${MAX_RETRIES} attempts.`);
                    }
                    return;
                }
                
                // 继续下一次循环重试
            }

        } catch (error) {
            // 只有发生网络断开、DNS 解析失败等底层错误时才会走到这里
            console.error(`[Proxy Fatal Error on attempt ${attempt}]: ${error.message}`);
            
            // 如果已达最大重试次数，返回 500 错误
            if (attempt >= MAX_RETRIES) {
                if (!res.headersSent) {
                    res.status(500).send(`Proxy Error: ${error.message}`);
                }
                return;
            }
            // 继续下一次循环重试
        }
    }
});

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