server.js

const express = require('express');
const { createProxyMiddleware, fixRequestBody } = require('http-proxy-middleware');
const fs = require('fs');
const path = require('path');

const app = express();
const proxyCache = {}; 
const stats = {}; 

// Load Endpoints & Initialize Stats
let endpoints = {};
try {
    endpoints = JSON.parse(fs.readFileSync(path.join(__dirname, 'endpoints.json'), 'utf8'));
    console.log("✔ Loaded Endpoints:", Object.keys(endpoints).join(', '));
    
    for (const key in endpoints) {
        stats[key] = { count: 0, totalTime: 0, avgTime: 0 };
    }
} catch (e) {
    console.error("✘ Error: endpoints.json missing or invalid.");
}

// 1. GLOBAL CORS INTERCEPTOR
app.use((req, res, next) => {
    res.setHeader('Access-Control-Allow-Origin', '*');
    res.setHeader('Access-Control-Allow-Methods', 'GET, POST, OPTIONS, PUT, PATCH, DELETE');
    res.setHeader('Access-Control-Allow-Headers', 'Content-Type, Authorization');
    
    if (req.method === 'OPTIONS') return res.sendStatus(200);
    next();
});

app.use(express.json({ limit: '50mb' }));

// 2. OPENAI OMNI MODERATION MIDDLEWARE
async function checkModeration(req, res, next) {
    if (req.method !== 'POST' || !req.body?.messages) return next();

    try {
        const textToCheck = req.body.messages.map(m => m.content).join('\n');
        if (!textToCheck.trim()) return next();

        if (!process.env.OPENAI_API_KEY) {
            console.warn("[Moderation] Skipped: OPENAI_API_KEY is not defined in Space Env Settings.");
            return next();
        }

        const modRes = await fetch('https://api.openai.com/v1/moderations', {
            method: 'POST',
            headers: {
                'Content-Type': 'application/json',
                'Authorization': `Bearer ${process.env.OPENAI_API_KEY}`
            },
            body: JSON.stringify({
                model: "omni-moderation-latest",
                input: textToCheck
            })
        });

        const modData = await modRes.json();
        if (modData?.results?.[0]) {
            const results = modData.results[0];
            if (results.categories['sexual/minors']) {
                console.warn("[Moderation] Blocked CSAM content.");
                return res.status(403).json({ error: "Content policy violation." });
            }
        }
    } catch (err) {
        console.error("[Moderation Error]", err.message);
    }
    next();
}

// Serve UI assets out of the explicit public folder
app.use(express.static(path.join(__dirname, 'public')));

app.get('/api/endpoints', (req, res) => res.json(endpoints));
app.get('/api/stats', (req, res) => res.json(stats)); 

app.use('/:nick', async (req, res, next) => {
    const nick = req.params.nick;

    if (!endpoints[nick]) {
        return res.status(404).json({ error: `Endpoint '${nick}' not found.` });
    }

    if (!proxyCache[nick]) {
        proxyCache[nick] = createProxyMiddleware({
            target: endpoints[nick].url,
            changeOrigin: true,
            pathRewrite: { [`^/${nick}`]: '' },
            on: {
                proxyReq: (proxyReq, req, res) => {
                    if (req.method === 'POST' && req.body) {
                        if (nick === 'nim') {
                            req.body.chat_template_kwargs = { enable_thinking: true, clear_thinking: false };
                            if (req.body.messages && Array.isArray(req.body.messages)) {
                                req.body.messages.forEach(msg => {
                                    if (msg.role === 'developer') msg.role = 'system';
                                });
                            }
                        }
                        
                        // Re-serializes the modified body correctly
                        fixRequestBody(proxyReq, req);
                    }
                },
                proxyRes: (proxyRes, req, res) => {
                    // 🔥 BYPASS HUGGINGFACE BUFFERING: Keeps streams alive
                    res.setHeader('X-Accel-Buffering', 'no');
                },
                error: (err, req, res) => {
                    console.error(`[Error: ${nick}]`, err.message);
                    if (!res.headersSent) res.status(502).json({ error: "Provider Offline or Timeout." });
                }
            }
        });
    }

    return proxyCache[nick](req, res, next);
});

const PORT = 7860;
app.listen(PORT, '0.0.0.0', () => console.log(`🚀 Hub active on port ${PORT}`));