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

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

app.set('trust proxy', 1);
app.use(cors());

// CRÍTICO: 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 ---
function logError(providerId, reason) {
    const timestamp = new Date().toISOString();
    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: "pollinations", 
        url: "https://text.pollinations.ai/openai", 
        modelsUrl: "https://text.pollinations.ai/models",
        imageUrl: "https://image.pollinations.ai/prompt" // Endpoint nativo de imágenes
    },
    { 
        id: "airforce", 
        url: "https://api.airforce/v1/chat/completions",
        imageUrl: "https://api.airforce/v1/images/generations" // Soporte de imágenes
    },
    { 
        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, "pollinations": 0, "airforce": 0, "ventarys-mirror": 0 };

const limiter = rateLimit({
    windowMs: 60 * 1000, 
    max: 25, 
    message: { error: { message: "Límite alcanzado. Espera 1 minuto entre mensajes.", 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(id) {
    if (!id) return false;
    const lowerId = id.toLowerCase();
    return IMAGE_KEYWORDS.some(kw => lowerId.includes(kw));
}

function isAudioModel(id, type) {
    if (type === 'audio' || type === 'music') return true;
    if (!id) return false;
    const lowerId = id.toLowerCase();
    return AUDIO_KEYWORDS.some(kw => lowerId.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
                    // Eliminamos los modelos de música de raíz para que no ensucien ninguna lista
                    .filter(model => !isAudioModel(model.id || model.name, model.type))
                    .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();
        // Filtrar modelos asegurando que NO sean de imágenes
        const textModels = allModels.filter(m => !isImageModel(m.id) && m.type !== 'image').map(m => {
            m.type = 'text'; // Forzar etiqueta
            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();
        // Filtrar modelos asegurando que SÍ sean de imágenes
        let imageModels = allModels.filter(m => isImageModel(m.id) || m.type === 'image').map(m => {
            m.type = 'image'; // Forzar etiqueta
            return m;
        });

        // Garantizar que la App siempre reconozca los modelos base aunque la API origen no los liste
        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';
    const availableProviders = isImage ? PROVIDERS.filter(p => p.imageUrl) : PROVIDERS;
    
    const startTime = Date.now();
    let selectedProvider = null;

    // Sistema de cola inteligente
    while (Date.now() - startTime < QUEUE_TIMEOUT) {
        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) break;
        await new Promise(r => setTimeout(r, 1500));
    }

    if (!selectedProvider) {
        logError("ProxyMain", "Saturación - Todas las APIs están ocupadas");
        return res.status(503).json({ error: { message: "Todas las APIs están ocupadas. Por favor, reintenta.", code: 503 } });
    }

    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;

        // Adaptador especializado para Pollinations Imágenes (Formato nativo GET)
        if (isImage && selectedProvider.id === "pollinations") {
            const prompt = req.body.prompt || "A random image";
            const seed = Math.floor(Math.random() * 10000000);
            
            // Extraer resolución del body (por defecto cuadrado si no se provee)
            let width = 1024, height = 1024;
            if (req.body.size) {
                const parts = req.body.size.split('x');
                if (parts.length === 2) { width = parseInt(parts[0]); height = parseInt(parts[1]); }
            }
            
            // Ajustar el modelo si el usuario pidió uno en específico
            const pollModel = req.body.model && req.body.model.includes('flux') ? 'flux' : 'dall-e-3';

            targetUrl = `${targetUrl}/${encodeURIComponent(prompt)}?seed=${seed}&width=${width}&height=${height}&model=${pollModel}&nologo=true`;
            reqMethod = "GET";
            reqBody = undefined;
            delete fetchHeaders["Content-Type"];
        }

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

        if (!response.ok) {
            logError(selectedProvider.id, `Respuesta HTTP ${response.status} - ${response.statusText}`);
        }

        // --- MANEJO EXCLUSIVO DE IMÁGENES (Estandarización a Base64 - Estilo Aqua AI) ---
        if (isImage) {
            const contentType = response.headers.get("content-type") || "";

            if (contentType.includes("application/json")) {
                const jsonResp = await response.json();
                
                // Si la API remota nos responde con JSON, extraemos las URL para transformarlas a Base64
                if (jsonResp.data && Array.isArray(jsonResp.data)) {
                    for (let item of jsonResp.data) {
                        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 convirtiendo URL a Base64: ${e.message}`);
                            }
                        }
                    }
                }
                releaseSlot();
                return res.status(response.status).json(jsonResp);
            } 
            else if (contentType.includes("image/")) {
                // Si la API nos devuelve el binario directo (Ej. Pollinations)
                const arrayBuffer = await response.arrayBuffer();
                const b64 = Buffer.from(arrayBuffer).toString('base64');
                releaseSlot();
                
                // Estandarizar respuesta para tu frontend (Aqua AI / OpenAI compatible)
                return res.status(200).json({
                    created: Math.floor(Date.now() / 1000),
                    data: [{ b64_json: b64 }]
                });
            } 
            else {
                // Caída libre para errores o comportamientos inesperados de los proveedores
                const textResp = await response.text();
                releaseSlot();
                return res.status(response.status).type(contentType).send(textResp);
            }
        } 
        
        // --- MANEJO EXCLUSIVO DE TEXTO (Streaming) ---
        res.writeHead(response.status, {
            'Content-Type': response.headers.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 roto a mitad de respuesta: ${err.message}`);
                releaseSlot();
            });
            req.on('close', releaseSlot); 
        } else {
            releaseSlot();
            res.end();
        }

    } catch (err) {
        releaseSlot();
        logError(selectedProvider ? selectedProvider.id : "Proxy", `Fallo o Timeout conectando al proveedor: ${err.message}`);
        res.status(500).json({ error: { message: "Error de conexión interna."} });
    }
});

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