import express from 'express';
import cors from 'cors';
import rateLimit from 'express-rate-limit';
import { Readable } from 'stream';

const app = express();
const PORT = 7860;

// --- CONFIGURACIÓN DE SEGURIDAD Y PRIVACIDAD ---
app.set('trust proxy', 1);
app.disable('x-powered-by'); // Oculta que estamos usando Express

app.use(cors());

// Inyección de cabeceras de seguridad estándar
app.use((req, res, next) => {
    res.setHeader('X-Content-Type-Options', 'nosniff');
    res.setHeader('X-Frame-Options', 'DENY');
    res.setHeader('X-XSS-Protection', '1; mode=block');
    res.setHeader('Strict-Transport-Security', 'max-age=31536000; includeSubDomains');
    res.setHeader('Referrer-Policy', 'no-referrer'); // Protege la privacidad del origen
    next();
});

// Aumentar límite a 50mb para soportar imágenes en Base64 (Visión / Multimodal)
app.use(express.json({ limit: '50mb' }));
app.use(express.urlencoded({ limit: '50mb', extended: true }));

// --- FUNCIÓN DE LOGS (Sanitizada) ---
function logError(providerId, reason) {
    const timestamp = new Date().toISOString();
    // No guardamos IPs ni prompts de usuarios, solo el estado de nuestros proveedores
    console.error(`[${timestamp}] [ERROR] Proveedor: ${providerId} | Motivo: ${reason}`);
}

// --- CONFIGURACIÓN DE PROVEEDORES ---
const PROVIDERS = [
    { 
        id: "llm7", 
        url: "https://api.llm7.io/v1/chat/completions" 
    },
    { 
        id: "airforce", 
        url: "https://api.airforce/v1/chat/completions",
        imageUrl: "https://api.airforce/v1/images/generations"
    },
    { 
        id: "ventarys-mirror", 
        url: "https://ventarys-mirror-1.hf.space/v1/chat/completions", 
        proxySecret: "sk-52650650a50f0v10vg150vs0v" 
    }
];

const MAX_PER_PROVIDER = 3; 
const QUEUE_TIMEOUT = 25000; 

let currentLoad = { "llm7": 0, "airforce": 0, "ventarys-mirror": 0 };

// --- RATE LIMITING (Basado en la IP real del usuario) ---
const limiter = rateLimit({
    windowMs: 60 * 1000, 
    max: 25, 
    keyGenerator: (req) => req.ip, 
    message: { error: { message: "Límite de solicitudes alcanzado. Por favor, espera un momento.", code: 429 } },
    standardHeaders: true,
    legacyHeaders: false,
});

// --- AYUDANTES PARA FILTRAR MODELOS ---
const IMAGE_KEYWORDS = ["flux", "dall", "midjourney", "sdxl", "stable-diffusion", "image", "vision"];
const AUDIO_KEYWORDS = ["suno", "udio", "music", "audio", "song", "voice", "tts"];

function isImageModel(model) {
    if (!model) return false;
    if (model.type === 'image' || model.supports_images === true) return true;
    
    const id = (model.id || model.name || "").toLowerCase();
    return IMAGE_KEYWORDS.some(kw => id.includes(kw));
}

function isAudioModel(model) {
    if (model.type === 'audio' || model.type === 'music') return true;
    const id = (model.id || model.name || "").toLowerCase();
    return AUDIO_KEYWORDS.some(kw => id.includes(kw));
}

async function fetchAllModels() {
    const fetchPromises = PROVIDERS.map(async (provider) => {
        const modelsUrl = provider.modelsUrl || provider.url.replace("/chat/completions", "/models");
        const fetchHeaders = { "Content-Type": "application/json" };
        
        if (provider.apiKey) fetchHeaders["Authorization"] = `Bearer ${provider.apiKey}`;
        if (provider.proxySecret) fetchHeaders["X-Proxy-Secret"] = provider.proxySecret;

        try {
            const resp = await fetch(modelsUrl, { method: "GET", headers: fetchHeaders });
            if (!resp.ok) return [];
            
            const json = await resp.json();
            let modelsArray = [];

            if (Array.isArray(json)) modelsArray = json; 
            else if (json && Array.isArray(json.data)) modelsArray = json.data; 

            if (modelsArray.length > 0) {
                return modelsArray
                    // 1. Filtrar modelos de audio/música
                    .filter(model => !isAudioModel(model))
                    // 2. CRÍTICO: Permitir estrictamente modelos con precio 0 (si la API reporta el precio)
                    .filter(model => {
                        if (model.pricepermilliontokens !== undefined && model.pricepermilliontokens !== null) {
                            return model.pricepermilliontokens === 0;
                        }
                        // Si el proveedor (ej. llm7) no envía el campo de precio, lo asumimos como válido
                        return true;
                    })
                    // 3. Procesar y formatear campos
                    .map(model => ({
                        ...model,
                        id: model.id || model.name,
                        owned_by: provider.id 
                    }));
            }
            return [];
        } catch (error) {
            return [];
        }
    });

    const results = await Promise.allSettled(fetchPromises);
    let allModels = [];
    results.forEach(result => {
        if (result.status === "fulfilled") allModels = allModels.concat(result.value);
    });
    return allModels;
}

// --- RUTAS INFORMATIVAS ---
app.get('/health', (req, res) => {
    res.json({ 
        status: "online", 
        type: "hf-node-proxy",
        providers: PROVIDERS.map(p => p.id),
        current_load: currentLoad 
    });
});

// Endpoint para modelos de TEXTO
app.get('/v1/models', async (req, res) => {
    try {
        const allModels = await fetchAllModels();
        
        const textModels = allModels
            .filter(m => !isImageModel(m) && m.supports_chat !== false)
            .map(m => {
                m.type = 'text';
                return m;
            });

        res.json({ object: "list", data: textModels });
    } catch (error) {
        logError("ProxyMain", "Error recuperando modelos de texto.");
        res.status(500).json({ error: "No se pudieron recuperar los modelos." });
    }
});

// Endpoint para modelos de IMAGEN
app.get('/v1/images/models', async (req, res) => {
    try {
        const allModels = await fetchAllModels();
        
        let imageModels = allModels
            .filter(m => isImageModel(m))
            .map(m => {
                m.type = 'image';
                return m;
            });

        const baseImages = [
            { id: "flux", object: "model", type: "image", owned_by: "system", tier: "standard" },
            { id: "dall-e-3", object: "model", type: "image", owned_by: "system", tier: "standard" }
        ];

        baseImages.forEach(baseMod => {
            if (!imageModels.find(m => m.id.toLowerCase() === baseMod.id)) {
                imageModels.push(baseMod);
            }
        });

        res.json({ object: "list", data: imageModels });
    } catch (error) {
        logError("ProxyMain", "Error recuperando modelos de imagen.");
        res.status(500).json({ error: "No se pudieron recuperar los modelos." });
    }
});

// --- RUTA PRINCIPAL DE GENERACIÓN ---
app.post(['/v1/chat/completions', '/v1/images/generations'], limiter, async (req, res) => {
    const isImage = req.path === '/v1/images/generations';
    
    let availableProviders = isImage ? PROVIDERS.filter(p => p.imageUrl) : [...PROVIDERS];
    const startTime = Date.now();
    let responseSent = false;

    while (availableProviders.length > 0 && Date.now() - startTime < QUEUE_TIMEOUT) {
        let selectedProvider = null;
        
        let shuffled = [...availableProviders].sort(() => Math.random() - 0.5);
        
        for (let provider of shuffled) {
            if (currentLoad[provider.id] < MAX_PER_PROVIDER) {
                selectedProvider = provider;
                currentLoad[provider.id]++;
                break;
            }
        }

        if (!selectedProvider) {
            await new Promise(r => setTimeout(r, 1500));
            continue;
        }

        let isReleased = false;
        const releaseSlot = () => {
            if (!isReleased) {
                currentLoad[selectedProvider.id] = Math.max(0, currentLoad[selectedProvider.id] - 1);
                isReleased = true;
            }
        };

        try {
            let targetUrl = isImage ? selectedProvider.imageUrl : selectedProvider.url;
            let reqMethod = "POST";
            let reqBody = JSON.stringify(req.body);
            const fetchHeaders = { "Content-Type": "application/json" };
            
            if (selectedProvider.apiKey) fetchHeaders["Authorization"] = `Bearer ${selectedProvider.apiKey}`;
            if (selectedProvider.proxySecret) fetchHeaders["X-Proxy-Secret"] = selectedProvider.proxySecret;

            const response = await fetch(targetUrl, {
                method: reqMethod,
                headers: fetchHeaders,
                body: reqBody
            });

            if (!response.ok) {
                logError(selectedProvider.id, `Fallo con código HTTP ${response.status}`);
                releaseSlot();
                availableProviders = availableProviders.filter(p => p.id !== selectedProvider.id);
                continue; 
            }

            // Sanitización: Evitar devolver cookies o cabeceras de servidor del proveedor original
            const responseHeaders = new Headers(response.headers);
            responseHeaders.delete('set-cookie');
            responseHeaders.delete('server');
            responseHeaders.delete('x-powered-by');
            responseHeaders.delete('cf-ray');

            if (isImage) {
                const contentType = responseHeaders.get("content-type") || "";

                if (contentType.includes("application/json")) {
                    const jsonResp = await response.json();
                    
                    let dataArray = jsonResp.data;
                    if (!dataArray && jsonResp.url) {
                        dataArray = [{ url: jsonResp.url }];
                    }
                    
                    if (dataArray && Array.isArray(dataArray)) {
                        for (let item of dataArray) {
                            if (item.url && !item.b64_json) {
                                try {
                                    const imgRes = await fetch(item.url);
                                    const arrayBuffer = await imgRes.arrayBuffer();
                                    item.b64_json = Buffer.from(arrayBuffer).toString('base64');
                                    delete item.url;
                                } catch (e) {
                                    logError(selectedProvider.id, `Fallo convitiendo URL a Base64.`);
                                }
                            }
                        }
                        releaseSlot();
                        responseSent = true;
                        return res.status(response.status).json({
                            created: Math.floor(Date.now() / 1000),
                            data: dataArray
                        });
                    }
                    
                    releaseSlot();
                    responseSent = true;
                    return res.status(response.status).json(jsonResp);
                } 
                else if (contentType.includes("image/")) {
                    const arrayBuffer = await response.arrayBuffer();
                    const b64 = Buffer.from(arrayBuffer).toString('base64');
                    releaseSlot();
                    
                    responseSent = true;
                    return res.status(200).json({
                        created: Math.floor(Date.now() / 1000),
                        data: [{ b64_json: b64 }]
                    });
                } 
                else {
                    const textResp = await response.text();
                    releaseSlot();
                    responseSent = true;
                    return res.status(response.status).type(contentType).send(textResp);
                }
            } 
            
            // Streaming (Texto)
            res.writeHead(response.status, {
                'Content-Type': responseHeaders.get('content-type') || 'text/event-stream',
                'Cache-Control': 'no-cache',
                'Connection': 'keep-alive'
            });

            if (response.body) {
                const stream = Readable.fromWeb(response.body);
                stream.pipe(res);
                
                stream.on('end', releaseSlot);
                stream.on('error', (err) => {
                    logError(selectedProvider.id, `Stream interrumpido.`);
                    releaseSlot();
                });
                req.on('close', releaseSlot); 
            } else {
                releaseSlot();
                res.end();
            }
            
            responseSent = true;
            return; 

        } catch (err) {
            logError(selectedProvider.id, `Excepción de red conectando al upstream.`);
            releaseSlot();
            availableProviders = availableProviders.filter(p => p.id !== selectedProvider.id);
        }
    }

    if (!responseSent) {
        logError("ProxyMain", "Todos los proveedores fallaron o se agotó el tiempo de espera.");
        return res.status(503).json({ error: { message: "El servicio no está disponible temporalmente. Inténtalo de nuevo.", code: 503 } });
    }
});

app.listen(PORT, '0.0.0.0', () => {
    console.log(`🚀 API Proxy (Node.js) corriendo seguro en el puerto ${PORT}`);
});