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)); 

// 3. THE SAFE PROXY SETUP
app.use('/:nick', checkModeration, (req, res, next) => {
    const { nick } = req.params;

    if (nick === 'api' || !endpoints[nick]) return next();

    // In-memory request timing metrics
    const startTime = Date.now();
    res.on('finish', () => {
        try {
            const duration = Date.now() - startTime;
            if (stats[nick]) {
                stats[nick].count++;
                stats[nick].totalTime += duration;
                stats[nick].avgTime = Math.round(stats[nick].totalTime / stats[nick].count);
            }
        } catch (e) {
            console.error("Metrics logging error:", e);
        }
    });

    if (!proxyCache[nick]) {
        const target = endpoints[nick].url;
        console.log(`[Setup] Routing /${nick} -> ${target}`);

        proxyCache[nick] = createProxyMiddleware({
            target: target,
            changeOrigin: true,
            pathRewrite: { [`^/${nick}`]: '' },
                        onProxyReq: (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
                    fixRequestBody(proxyReq, req);
                    
                    // 🔥 THE FIX: Force the proxy request to finish!
                    // Since the async middleware drained the original stream, 
                    // we have to manually tell the proxy we are done writing.
                    proxyReq.end(); 
                }
            },

          
            onError: (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}`));