server.js

import express from 'express';
import axios from 'axios';
import { Client } from '@gradio/client';
import dotenv from 'dotenv';

dotenv.config();

// --- Hugging Face Space Auto-Restarter ---
const HF_TOKEN = process.env.HF_TOKEN;
const SPACE_ID = process.env.HF_SPACE_ID || 'shashwatIDR/vision';
const HF_API_BASE = 'https://huggingface.co/api';

async function restartSpace() {
  console.log(`[🔄 Attempting restart] ${SPACE_ID}`);
  try {
    const response = await axios.post(`${HF_API_BASE}/spaces/${SPACE_ID}/restart`, {}, {
      headers: {
        Authorization: `Bearer ${HF_TOKEN}`,
        'Content-Type': 'application/json'
      },
      timeout: 30000
    });
    console.log(`[✅ Restarted] ${SPACE_ID} at ${new Date().toLocaleString()}`);
    if (response.data) {
      console.log('[📊 Response]', response.data);
    }
  } catch (err) {
    if (err.code === 'ENOTFOUND') {
      console.error('[❌ DNS Error] Cannot resolve huggingface.co - check your internet connection');
    } else if (err.code === 'ETIMEDOUT') {
      console.error('[❌ Timeout] Request timed out - try again later');
    } else if (err.response) {
      console.error(`[❌ HTTP ${err.response.status}]`, err.response.data || err.response.statusText);
      if (err.response.status === 401) {
        console.error('[🔑 Auth Error] Check your Hugging Face token');
      } else if (err.response.status === 404) {
        console.error('[🔍 Not Found] Space may not exist or token lacks permissions');
      }
    } else {
      console.error('[❌ Failed to Restart]', err.message);
    }
  }
}

async function testHFAPI() {
  try {
    console.log('[🔍 Testing Hugging Face API endpoint...]');
    const response = await axios.get(`${HF_API_BASE}/spaces`, {
      headers: { Authorization: `Bearer ${HF_TOKEN}` },
      timeout: 10000,
      params: { limit: 1 }
    });
    console.log('[✅ API Test] Endpoint is working');
    return true;
  } catch (err) {
    console.error('[❌ API Test Failed]', err.response?.status || err.code || err.message);
    return false;
  }
}

// On server startup, test API and schedule restarts
(async () => {
  if (!HF_TOKEN) {
    console.error('[❌ No HF_TOKEN] Hugging Face token not set. Skipping auto-restart scheduling.');
    return;
  }
  console.log('[🚀 Starting] HuggingFace Space Auto-Restarter');
  const apiWorking = await testHFAPI();
  if (!apiWorking) {
    console.log('[⚠️ Warning] API test failed, but proceeding anyway...');
  }
  // Only schedule periodic restarts (no immediate restart on startup)
  setInterval(restartSpace, 5 * 60 * 1000);
  console.log('[⏰ Scheduled] Space will restart every 5 minutes');
})();

const app = express();
const port = process.env.PORT || 3000;
app.use(express.json());

// Serve static files from public
app.use(express.static('public'));

// Serve index.html at root
app.get('/', (req, res) => {
  res.sendFile(__dirname + '/public/index.html');
});

// Image Generation endpoint (Heartsync)
app.post('/api/generate', async (req, res) => {
    res.setHeader('Content-Type', 'text/event-stream');
    res.setHeader('Cache-Control', 'no-cache');
    res.setHeader('Connection', 'keep-alive');

    try {
        const { prompt } = req.body;
        const sessionHash = Math.random().toString(36).substring(2, 15);

        // First, join the queue
        const joinResponse = await axios.post('https://heartsync-nsfw-uncensored.hf.space/gradio_api/queue/join', {
            data: [
                prompt,
                "text, talk bubble, low quality, watermark, signature",
                0,
                true,
                1024,
                1024,
                7,
                28
            ],
            event_data: null,
            fn_index: 2,
            trigger_id: 14,
            session_hash: sessionHash
        }, {
            headers: {
                'Content-Type': 'application/json',
                'Referer': 'https://heartsync-nsfw-uncensored.hf.space/'
            }
        });

        // Poll for results
        let attempts = 0;
        const maxAttempts = 60; // Increased attempts for potentially longer queues
        
        while (attempts < maxAttempts) {
            const dataResponse = await axios.get(
                `https://heartsync-nsfw-uncensored.hf.space/gradio_api/queue/data?session_hash=${encodeURIComponent(sessionHash)}`,
                {
                    headers: {
                        'Accept': 'text/event-stream',
                        'Referer': 'https://heartsync-nsfw-uncensored.hf.space/'
                    }
                }
            );

            const data = dataResponse.data;
            
            // Manually parse the event stream data
            const lines = data.split('\n');
            let eventData = '';
            let foundEvent = false;
            for (const line of lines) {
                if (line.startsWith('data: ')) {
                    eventData += line.substring(6);
                } else if (line === '') {
                    // Process eventData when an empty line is encountered
                    if (eventData) {
                        try {
                            const json = JSON.parse(eventData);
                            console.log(`[Image] Parsed event: ${json.msg}`);
                            
                            if (json.msg === 'process_completed' && json.output?.data?.[0]?.url) {
                                const imageUrl = json.output.data[0].url;
                                
                                // Send the original image URL back to the client as an SSE
                                res.write(`data: ${JSON.stringify({
                                    type: 'success',
                                    originalUrl: imageUrl
                                })}\n\n`);
                                res.end(); // End the connection after sending the final event
                                return; // Exit the function
                            }

                            if (json.msg === 'estimation') {
                                res.write(`data: ${JSON.stringify({ type: 'estimation', queueSize: json.queue_size, eta: json.rank_eta })}\n\n`);
                                foundEvent = true;
                            } else if (json.msg === 'process_starts') {
                                res.write(`data: ${JSON.stringify({ type: 'processing' })}\n\n`);
                                foundEvent = true;
                            } else if (json.msg === 'process_failed') {
                                console.log(`[Image] Process failed:`, json);
                                res.write(`data: ${JSON.stringify({ type: 'error', message: 'Generation failed on server' })}\n\n`);
                                res.end();
                                return;
                            }

                        } catch (error) {
                            console.error('[Image] Error parsing event data:', error, 'Raw data:', eventData);
                        }
                        eventData = ''; // Reset for the next event
                    }
                } else if (line.startsWith(':')) {
                    // Ignore comment lines
                    continue;
                } else if (line !== '') {
                     // If line is not data or comment, it might be part of a multi-line data chunk, append to eventData
                     eventData += line;
                }
            }

            // If loop finishes without 'process_completed', wait and try again
            await new Promise(resolve => setTimeout(resolve, 2000)); // Wait longer before next poll attempt
            attempts++;
        }

        // If the loop times out without success
        if (!foundEvent) {
            console.log(`[Image] No events found in response, raw data:`, data.substring(0, 500));
        }
        res.status(500).write(`data: ${JSON.stringify({ type: 'error', message: 'Generation timed out' })}\n\n`);
        res.end();

    } catch (error) {
        console.error('Generation endpoint error:', error);
        res.status(500).write(`data: ${JSON.stringify({ type: 'error', message: error.message || 'Failed to generate image' })}\n\n`);
        res.end();
    }
});

// Flux model generation endpoint using black-forest-labs/FLUX.1-dev Gradio client
app.post('/api/generate/flux', async (req, res) => {
    res.setHeader('Content-Type', 'text/event-stream');
    res.setHeader('Cache-Control', 'no-cache');
    res.setHeader('Connection', 'keep-alive');
    try {
        const {
            prompt = '',
            width = 1024,
            height = 1024,
            guidance_scale = 3.5,
            num_inference_steps = 28,
            seed = 0,
            randomize_seed = true
        } = req.body;

        if (!prompt) {
            res.write(`data: ${JSON.stringify({ type: 'error', message: 'Prompt is required' })}\n\n`);
            res.end();
            return;
        }

        // Simulate estimation and processing events for UI feedback
        res.write(`data: ${JSON.stringify({ type: 'estimation', queueSize: 1, eta: 5 })}\n\n`);
        await new Promise(resolve => setTimeout(resolve, 1000));
        res.write(`data: ${JSON.stringify({ type: 'processing' })}\n\n`);

        // Connect to FLUX.1-dev Gradio client
        const client = await Client.connect("black-forest-labs/FLUX.1-dev");
        const result = await client.predict("/infer", {
            prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps
        });

        // result.data: [imageObj, seed]
        const imageObj = result.data[0];
        const imageUrl = (imageObj && imageObj.url) ? imageObj.url : (typeof imageObj === 'string' ? imageObj : null);
        if (imageUrl) {
            res.write(`data: ${JSON.stringify({ type: 'success', imageUrl, seed: result.data[1] })}\n\n`);
        } else {
            res.write(`data: ${JSON.stringify({ type: 'error', message: 'No image URL returned' })}\n\n`);
        }
        res.end();
    } catch (err) {
        res.write(`data: ${JSON.stringify({ type: 'error', message: err.message })}\n\n`);
        res.end();
    }
});

app.listen(port, () => {
  console.log(`Server running on port ${port}`);
});