server.js

import express from 'express';
import cors from 'cors';
import multer from 'multer';
import axios from 'axios';
import { Client } from '@gradio/client';
import dotenv from 'dotenv';
import bcrypt from 'bcrypt';
import jwt from 'jsonwebtoken';
import FormData from 'form-data';
import pg from 'pg';
import fetch from 'node-fetch';

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 { Pool } = pg;

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

// Middleware
app.use(cors({ origin: '*', credentials: true }));
app.use(express.json());
app.use(express.static('public'));

// Database connection
const pool = new Pool({
  connectionString: 'postgresql://img_owner:npg_PgXtYG4py8Vb@ep-wandering-salad-a1eva4cg-pooler.ap-southeast-1.aws.neon.tech/img?sslmode=require',
  ssl: {
    rejectUnauthorized: false
  }
});

// Create tables if they don't exist
const initializeDatabase = async () => {
  try {
    await pool.query(`
      CREATE TABLE IF NOT EXISTS users (
        id SERIAL PRIMARY KEY,
        username VARCHAR(255) NOT NULL UNIQUE,
        email VARCHAR(255) NOT NULL UNIQUE,
        password VARCHAR(255) NOT NULL,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
      );

      CREATE TABLE IF NOT EXISTS images (
        id SERIAL PRIMARY KEY,
        user_id INTEGER REFERENCES users(id),
        image_url VARCHAR(255) NOT NULL,
        uploaded_url VARCHAR(255),
        prompt TEXT NOT NULL,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
      );

      CREATE TABLE IF NOT EXISTS user_generations (
        id SERIAL PRIMARY KEY,
        user_id INTEGER REFERENCES users(id),
        original_url VARCHAR(255) NOT NULL,
        uploaded_url VARCHAR(255),
        prompt TEXT NOT NULL,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        deleted_at TIMESTAMP
      );
    `);
    // Add censored column if it doesn't exist
    await pool.query(`ALTER TABLE images ADD COLUMN IF NOT EXISTS censored BOOLEAN DEFAULT false;`);
    console.log('Database initialized successfully');
  } catch (error) {
    console.error('Error initializing database:', error);
    throw error;
  }
};

initializeDatabase();

// Authentication middleware
const authenticateToken = (req, res, next) => {
  const authHeader = req.headers['authorization'];
  const token = authHeader && authHeader.split(' ')[1];

  if (!token) return res.sendStatus(401);

  jwt.verify(token, process.env.JWT_SECRET || 'your-secret-key', (err, user) => {
    if (err) return res.sendStatus(403);
    req.user = user;
    next();
  });
};

// Gemini censorship check function
const GEMINI_API_KEY = 'AIzaSyCq23lcvpPfig6ifq1rmt-z11vKpMvDD4I'; // <-- Updated Gemini API key
async function checkPromptCensorship(prompt) {
  const apiKey = GEMINI_API_KEY;
  const url = `https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key=${apiKey}`;
  const systemInstruction =
    'You are a content moderation AI. If the following prompt contains any NSFW, adult, explicit, sexual, violent, or otherwise censored content, respond ONLY with the word "true". If it is safe and contains no censored elements, respond ONLY with the word "false". Do not explain.';
  try {
    const response = await axios.post(url, {
      contents: [
        {
          parts: [
            { text: `${systemInstruction}\nPrompt: ${prompt}` }
          ]
        }
      ]
    });
    const text = response.data?.candidates?.[0]?.content?.parts?.[0]?.text?.trim().toLowerCase();
    if (text === 'true') return true;
    if (text === 'false') return false;
    // fallback: treat as not censored if unclear
    return false;
  } catch (e) {
    console.error('Gemini censorship API error:', e.message);
    // fallback: treat as not censored if error
    return false;
  }
}

// Routes
app.post('/api/register', async (req, res) => {
  try {
    const { username, email, password } = req.body;
    
    if (!username || !email || !password) {
      return res.status(400).json({ error: 'All fields are required' });
    }
    
    // Check if email already exists
    const emailCheck = await pool.query('SELECT id FROM users WHERE email = $1', [email]);
    if (emailCheck.rows.length > 0) {
      return res.status(400).json({ error: 'Email already registered' });
    }
    
    // Check if username already exists
    const usernameCheck = await pool.query('SELECT id FROM users WHERE username = $1', [username]);
    if (usernameCheck.rows.length > 0) {
      return res.status(400).json({ error: 'Username already taken' });
    }
    
    const hashedPassword = await bcrypt.hash(password, 10);
    
    const result = await pool.query(
      'INSERT INTO users (username, email, password) VALUES ($1, $2, $3) RETURNING id, username, email',
      [username, email, hashedPassword]
    );
    
    const token = jwt.sign({ id: result.rows[0].id }, process.env.JWT_SECRET || 'your-secret-key');
    res.json({ 
      token, 
      user: result.rows[0] 
    });
  } catch (error) {
    console.error('Registration error:', error);
    res.status(500).json({ error: 'Error during registration' });
  }
});

app.post('/api/login', async (req, res) => {
  try {
    const { email, password } = req.body;
    
    if (!email || !password) {
      return res.status(400).json({ error: 'Email and password are required' });
    }
    
    const result = await pool.query('SELECT * FROM users WHERE email = $1', [email]);
    
    if (result.rows.length === 0) {
      return res.status(401).json({ error: 'Invalid email or password' });
    }

    const user = result.rows[0];
    const validPassword = await bcrypt.compare(password, user.password);
    
    if (!validPassword) {
      return res.status(401).json({ error: 'Invalid email or password' });
    }

    const token = jwt.sign({ id: user.id }, process.env.JWT_SECRET || 'your-secret-key');
    res.json({ 
      token, 
      user: { 
        id: user.id, 
        username: user.username, 
        email: user.email 
      } 
    });
  } catch (error) {
    console.error('Login error:', error);
    res.status(500).json({ error: 'Error during login' });
  }
});

app.get('/api/reset-db', async (req, res) => {
  try {
    await pool.query(`
      DROP TABLE IF EXISTS user_generations;
      DROP TABLE IF EXISTS images;
      DROP TABLE IF EXISTS users;
    `);
    await initializeDatabase();
    res.json({ success: true, message: 'Database reset successfully' });
  } catch (error) {
    console.error('Error resetting database:', error);
    res.status(500).json({ error: 'Failed to reset database' });
  }
});

app.post('/api/upload-image', authenticateToken, async (req, res) => {
  try {
    const { imageUrl, prompt } = req.body;
    const userId = req.user.id;

    // First, download the image from Heartsync
    const imageResponse = await axios.get(imageUrl, { responseType: 'arraybuffer' });
    const imageBuffer = Buffer.from(imageResponse.data);

    // Create form data for File2Link
    const formData = new FormData();
    formData.append('file', new Blob([imageBuffer]), 'image.jpg');

    // Upload to File2Link
    const uploadResponse = await axios.post('https://file2link-ol4p.onrender.com/upload', formData, {
      headers: {
        ...formData.getHeaders()
      }
    });

    const uploadedUrl = uploadResponse.data.access_url;

    // Check censorship
    const censored = await checkPromptCensorship(prompt);

    // Save to images table for community gallery
    const imageResult = await pool.query(
      'INSERT INTO images (user_id, image_url, uploaded_url, prompt, censored) VALUES ($1, $2, $3, $4, $5) RETURNING *',
      [userId, imageUrl, uploadedUrl, prompt, censored]
    );

    // Save to user_generations table for user history
    const generationResult = await pool.query(
      'INSERT INTO user_generations (user_id, original_url, uploaded_url, prompt) VALUES ($1, $2, $3, $4) RETURNING *',
      [userId, imageUrl, uploadedUrl, prompt]
    );

    res.json({ 
      success: true, 
      image: imageResult.rows[0],
      generation: generationResult.rows[0]
    });
  } catch (error) {
    console.error('Error uploading image:', error);
    res.status(500).json({ error: 'Error uploading image' });
  }
});

app.get('/api/community-images', async (req, res) => {
  try {
    const result = await pool.query(`
      SELECT i.id, i.image_url, i.uploaded_url, i.prompt, i.created_at, i.censored, u.username 
      FROM images i 
      JOIN users u ON i.user_id = u.id 
      ORDER BY i.created_at DESC
    `);
    res.json(result.rows);
  } catch (error) {
    console.error('Error fetching community images:', error);
    res.status(500).json({ error: 'Error fetching community images' });
  }
});

app.get('/api/user-generations', authenticateToken, async (req, res) => {
  try {
    const userId = req.user.id;
    const result = await pool.query(
      'SELECT * FROM user_generations WHERE user_id = $1 AND deleted_at IS NULL ORDER BY created_at DESC',
      [userId]
    );
    res.json(result.rows);
  } catch (error) {
    console.error('Error fetching user generations:', error);
    res.status(500).json({ error: 'Error fetching user generations' });
  }
});

app.delete('/api/user-generations/:id', authenticateToken, async (req, res) => {
  try {
    const { id } = req.params;
    const userId = req.user.id;

    // First, get the image details to find the corresponding community image
    const generationResult = await pool.query(
      'SELECT original_url, uploaded_url FROM user_generations WHERE id = $1 AND user_id = $2',
      [id, userId]
    );

    if (generationResult.rows.length === 0) {
      return res.status(404).json({ error: 'Image not found' });
    }

    const { original_url, uploaded_url } = generationResult.rows[0];

    // Soft delete from user_generations
    await pool.query(
      'UPDATE user_generations SET deleted_at = CURRENT_TIMESTAMP WHERE id = $1 AND user_id = $2',
      [id, userId]
    );

    // Remove from community images (images table) - fix column name matching
    await pool.query(
      'DELETE FROM images WHERE user_id = $1 AND (image_url = $2 OR uploaded_url = $3)',
      [userId, original_url, uploaded_url]
    );

    res.json({ success: true });
  } catch (error) {
    console.error('Error deleting generation:', error);
    res.status(500).json({ error: 'Error deleting generation' });
  }
});

// Image Generation endpoint (backend only handles Heartsync interaction)
app.post('/api/generate', authenticateToken, 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();
    }
});

// Anime-UC Generation endpoint (using Heartsync API like original)
app.post('/api/generate/anime-uc', authenticateToken, 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);

        console.log(`[Anime-UC] Starting generation for prompt: "${prompt}" with session: ${sessionHash}`);

        // First, join the queue using Heartsync API
        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/'
            }
        });

        console.log(`[Anime-UC] Queue joined, event_id: ${joinResponse.data?.event_id}`);

        // Poll for results
        let attempts = 0;
        const maxAttempts = 60; // Heartsync typically faster
        
        while (attempts < maxAttempts) {
            try {
                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/'
                        },
                        timeout: 10000 // 10 second timeout
                    }
                );

                const data = dataResponse.data;
                console.log(`[Anime-UC] Poll attempt ${attempts + 1}, data length: ${data.length}`);
                
                // 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(`[Anime-UC] Parsed event: ${json.msg}`);
                                
                                if (json.msg === 'process_completed' && json.output?.data?.[0]?.url) {
                                    const imageUrl = json.output.data[0].url;
                                    console.log(`[Anime-UC] Found image:`, imageUrl);
                                    
                                    // Send the image URL back to the client as an SSE
                                    res.write(`data: ${JSON.stringify({
                                        type: 'success',
                                        originalUrl: imageUrl,
                                        allImages: [imageUrl]
                                    })}\n\n`);
                                    res.end();
                                    return;
                                }

                                if (json.msg === 'estimation') {
                                    res.write(`data: ${JSON.stringify({ type: 'estimation', queueSize: json.queue_size || 0, eta: json.rank_eta || null })}\n\n`);
                                    foundEvent = true;
                                } else if (json.msg === 'process_starts') {
                                    res.write(`data: ${JSON.stringify({ type: 'processing' })}\n\n`);
                                    foundEvent = true;
                                }

                            } catch (error) {
                                console.error('[Anime-UC] Error parsing event data:', error, 'Raw data:', eventData);
                            }
                            eventData = '';
                        }
                    } else if (line.startsWith(':')) {
                        continue;
                    } else if (line !== '') {
                         eventData += line;
                    }
                }

                // If no events were found, log the raw data for debugging
                if (!foundEvent && data.trim()) {
                    console.log(`[Anime-UC] No events found in response, raw data:`, data.substring(0, 500));
                }

            } catch (pollError) {
                console.error(`[Anime-UC] Poll error on attempt ${attempts + 1}:`, pollError.message);
            }

            await new Promise(resolve => setTimeout(resolve, 2000)); // Wait 2 seconds between polls
            attempts++;
        }

        console.log(`[Anime-UC] Generation timed out after ${maxAttempts} attempts`);
        res.status(500).write(`data: ${JSON.stringify({ type: 'error', message: 'Generation timed out' })}\n\n`);
        res.end();

    } catch (error) {
        console.error('[Anime-UC] 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 Generation endpoint (UNO-FLUX direct API with native fetch SSE polling)
app.get('/api/generate/flux', authenticateToken, async (req, res) => {
    try {
        const { prompt, aspect } = req.query;
        if (!prompt) {
            console.error('[Flux] No prompt provided');
            return res.status(400).json({ success: false, error: 'Prompt is required' });
        }

        // Aspect ratio logic
        let width = 1280, height = 720;
        if (aspect === '9:16') {
            width = 720; height = 1280;
        } else if (aspect === '1:1') {
            width = 1024; height = 1024;
        }

        const sessionHash = Math.random().toString(36).substring(2, 15);
        console.log(`[Flux] Starting UNO-FLUX generation for prompt: "${prompt}" | aspect: ${aspect} | width: ${width} | height: ${height} | session: ${sessionHash}`);

        // Join the queue using UNO-FLUX API
        const joinPayload = {
            data: [
                prompt,
                width,
                height,
                4,      // guidance
                25,     // num_steps
                -1,     // seed
                null,   // image_prompt1
                null,   // image_prompt2
                null,   // image_prompt3
                null    // image_prompt4
            ],
            event_data: null,
            fn_index: 0,
            trigger_id: 25,
            session_hash: sessionHash
        };
        console.log('[Flux] Sending queue/join payload:', JSON.stringify(joinPayload));
        let joinResponse;
        try {
            joinResponse = await fetch(
                'https://bytedance-research-uno-flux.hf.space/gradio_api/queue/join?__theme=system',
                {
                    method: 'POST',
                    headers: {
                        'Content-Type': 'application/json',
                        'Referer': 'https://bytedance-research-uno-flux.hf.space/',
                        'Origin': 'https://bytedance-research-uno-flux.hf.space'
                    },
                    body: JSON.stringify(joinPayload)
                }
            );
            const joinData = await joinResponse.json();
            console.log('[Flux] queue/join response:', joinData);
        } catch (err) {
            console.error('[Flux] Error joining UNO-FLUX queue:', err);
            return res.status(500).json({ success: false, error: 'Failed to join UNO-FLUX queue' });
        }

        // Poll for results using native fetch and streaming
        const pollUrl = `https://bytedance-research-uno-flux.hf.space/gradio_api/queue/data?session_hash=${encodeURIComponent(sessionHash)}`;
        console.log(`[Flux] Polling for results at: ${pollUrl}`);
        let found = false;
        let attempts = 0;
        const maxAttempts = 60;
        let buffer = '';
        while (attempts < maxAttempts && !found) {
            attempts++;
            try {
                const response = await fetch(pollUrl, {
                    headers: {
                        'Accept': 'text/event-stream',
                        'Referer': 'https://bytedance-research-uno-flux.hf.space/',
                        'Origin': 'https://bytedance-research-uno-flux.hf.space'
                    }
                });
                if (!response.ok) {
                    console.error(`[Flux] Poll error on attempt ${attempts}:`, response.status, await response.text());
                    await new Promise(resolve => setTimeout(resolve, 2000));
                    continue;
                }
                // Use Node.js stream reading
                for await (const chunk of response.body) {
                    buffer += chunk.toString('utf8');
                    const lines = buffer.split('\n');
                    for (let i = 0; i < lines.length - 1; i++) {
                        const line = lines[i];
                        if (line.startsWith('data: ')) {
                            try {
                                const json = JSON.parse(line.slice(6));
                                console.log('[Flux] Event:', json);
                                if (json.msg === 'process_completed' && json.output?.data?.[0]?.url) {
                                    const imageUrl = json.output.data[0].url;
                                    console.log(`[Flux] Success! Image URL: ${imageUrl}`);
                                    found = true;
                                    return res.json({ success: true, imageUrl });
                                }
                            } catch (e) {
                                console.error('[Flux] Parse error:', e, line);
                            }
                        }
                    }
                    buffer = lines[lines.length - 1]; // keep incomplete line
                }
            } catch (pollError) {
                console.error(`[Flux] Poll error on attempt ${attempts}:`, pollError);
            }
            await new Promise(resolve => setTimeout(resolve, 2000));
        }
        if (!found) {
            console.error(`[Flux] Generation timed out after ${maxAttempts} attempts`);
            res.status(500).json({ success: false, error: 'Generation timed out' });
        }
    } catch (error) {
        console.error('[Flux] Fatal error in endpoint:', error);
        res.status(500).json({ success: false, error: error.message || 'Failed to generate image' });
    }
});

// New endpoint to add community image (handles data from client-side upload)
app.post('/api/add-community-image', authenticateToken, async (req, res) => {
    try {
        const { originalUrl, uploadedUrl, prompt } = req.body;
        const userId = req.user.id;

        if (!originalUrl || !uploadedUrl || !prompt) {
             return res.status(400).json({ error: 'Missing required fields' });
        }

        // Check censorship
        const censored = await checkPromptCensorship(prompt);

        // Save to images table for community gallery
        const imageResult = await pool.query(
            'INSERT INTO images (user_id, image_url, uploaded_url, prompt, censored) VALUES ($1, $2, $3, $4, $5) RETURNING *',
            [userId, originalUrl, uploadedUrl, prompt, censored]
        );

        // Save to user_generations table for user history
        const generationResult = await pool.query(
            'INSERT INTO user_generations (user_id, original_url, uploaded_url, prompt) VALUES ($1, $2, $3, $4) RETURNING *',
            [userId, originalUrl, uploadedUrl, prompt]
        );

        res.json({
            success: true,
            image: imageResult.rows[0],
            generation: generationResult.rows[0]
        });
    } catch (error) {
        console.error('Error adding community image:', error);
        res.status(500).json({ error: 'Error adding image to community' });
    }
});

// Auth status endpoint
app.get('/api/auth/status', authenticateToken, async (req, res) => {
    try {
        const result = await pool.query(
            'SELECT id, username, email FROM users WHERE id = $1',
            [req.user.id]
        );
        
        if (result.rows.length === 0) {
            return res.status(401).json({ error: 'User not found' });
        }
        
        res.json({ user: result.rows[0] });
    } catch (error) {
        console.error('Auth status error:', error);
        res.status(500).json({ error: 'Error checking auth status' });
    }
});

// Admin endpoints (no authentication required for admin functions)

// Get image statistics
app.get('/api/admin/stats', async (req, res) => {
  try {
    const totalImages = await pool.query('SELECT COUNT(*) as count FROM images');
    const uniqueUsers = await pool.query('SELECT COUNT(DISTINCT user_id) as count FROM images');
    const last24Hours = await pool.query(`
      SELECT COUNT(*) as count FROM images 
      WHERE created_at >= NOW() - INTERVAL '24 hours'
    `);
    
    res.json({
      totalImages: parseInt(totalImages.rows[0].count),
      uniqueUsers: parseInt(uniqueUsers.rows[0].count),
      last24Hours: parseInt(last24Hours.rows[0].count)
    });
  } catch (error) {
    console.error('Error getting admin stats:', error);
    res.status(500).json({ error: 'Error getting statistics' });
  }
});

// Delete single image (admin)
app.delete('/api/admin/images/:id', async (req, res) => {
  try {
    const { id } = req.params;
    
    // Delete from images table
    const result = await pool.query('DELETE FROM images WHERE id = $1 RETURNING *', [id]);
    
    if (result.rows.length === 0) {
      return res.status(404).json({ error: 'Image not found' });
    }
    
    // Also delete from user_generations if it exists
    await pool.query(
      'DELETE FROM user_generations WHERE original_url = $1 OR uploaded_url = $2',
      [result.rows[0].image_url, result.rows[0].uploaded_url]
    );
    
    res.json({ success: true, message: 'Image deleted successfully' });
  } catch (error) {
    console.error('Error deleting image:', error);
    res.status(500).json({ error: 'Error deleting image' });
  }
});

// Delete all images (admin)
app.delete('/api/admin/images', async (req, res) => {
  try {
    // Delete all images from images table
    await pool.query('DELETE FROM images');
    
    // Clear user_generations table (soft delete all)
    await pool.query('UPDATE user_generations SET deleted_at = CURRENT_TIMESTAMP WHERE deleted_at IS NULL');
    
    res.json({ success: true, message: 'All images deleted successfully' });
  } catch (error) {
    console.error('Error deleting all images:', error);
    res.status(500).json({ error: 'Error deleting all images' });
  }
});

// Export all images data (admin)
app.get('/api/admin/export', async (req, res) => {
  try {
    const result = await pool.query(`
      SELECT i.id, i.image_url, i.uploaded_url, i.prompt, i.created_at, i.censored, u.username 
      FROM images i 
      JOIN users u ON i.user_id = u.id 
      ORDER BY i.created_at DESC
    `);
    const exportData = {
      images: result.rows,
      exportDate: new Date().toISOString(),
      totalImages: result.rows.length,
      uniqueUsers: new Set(result.rows.map(img => img.username)).size
    };
    res.setHeader('Content-Type', 'application/json');
    res.setHeader('Content-Disposition', `attachment; filename="community-images-${new Date().toISOString().split('T')[0]}.json"`);
    res.json(exportData);
  } catch (error) {
    console.error('Error exporting data:', error);
    res.status(500).json({ error: 'Error exporting data' });
  }
});

// Bulk update images (admin, no auth)
app.post('/api/admin/bulk-update-images', async (req, res) => {
  try {
    const images = req.body;
    if (!Array.isArray(images)) {
      return res.status(400).json({ error: 'Request body must be an array of images' });
    }
    let updated = 0, inserted = 0, errors = [];
    for (const img of images) {
      if (!img.id || !img.image_url || !img.uploaded_url) {
        errors.push({ id: img.id, error: 'Missing id, image_url, or uploaded_url' });
        continue;
      }
      try {
        // Try to update
        let updateFields = ['image_url = $1', 'uploaded_url = $2'];
        let updateValues = [img.image_url, img.uploaded_url, img.id];
        let paramIdx = 3;
        if (typeof img.prompt === 'string') {
          updateFields.push(`prompt = $${paramIdx}`);
          updateValues.splice(paramIdx - 1, 0, img.prompt);
          paramIdx++;
        }
        if (typeof img.censored === 'boolean') {
          updateFields.push(`censored = $${paramIdx}`);
          updateValues.splice(paramIdx - 1, 0, img.censored);
          paramIdx++;
        }
        const updateResult = await pool.query(
          `UPDATE images SET ${updateFields.join(', ')} WHERE id = $${paramIdx} RETURNING *`,
          updateValues
        );
        if (updateResult.rowCount > 0) {
          updated++;
        } else {
          // Insert if not found
          const insertFields = ['id', 'image_url', 'uploaded_url'];
          const insertValues = [img.id, img.image_url, img.uploaded_url];
          const insertParams = ['$1', '$2', '$3'];
          let insertIdx = 4;
          if (typeof img.prompt === 'string') {
            insertFields.push('prompt');
            insertValues.push(img.prompt);
            insertParams.push(`$${insertIdx++}`);
          }
          if (typeof img.censored === 'boolean') {
            insertFields.push('censored');
            insertValues.push(img.censored);
            insertParams.push(`$${insertIdx++}`);
          }
          await pool.query(
            `INSERT INTO images (${insertFields.join(', ')}) VALUES (${insertParams.join(', ')})`,
            insertValues
          );
          inserted++;
        }
      } catch (e) {
        errors.push({ id: img.id, error: e.message });
      }
    }
    res.json({ updated, inserted, errors });
  } catch (error) {
    console.error('Bulk update error:', error);
    res.status(500).json({ error: 'Bulk update failed' });
  }
});

// Admin endpoint to process all old images for censorship
app.post('/api/admin/process-old', async (req, res) => {
  try {
    // Get all images
    const result = await pool.query('SELECT id, prompt FROM images');
    let updated = 0, errors = [];
    for (const img of result.rows) {
      try {
        const censored = await checkPromptCensorship(img.prompt);
        await pool.query('UPDATE images SET censored = $1 WHERE id = $2', [censored, img.id]);
        updated++;
      } catch (e) {
        errors.push({ id: img.id, error: e.message });
      }
    }
    res.json({ updated, errors });
  } catch (error) {
    console.error('Process-old error:', error);
    res.status(500).json({ error: 'Failed to process old images' });
  }
});

// Proxy route for HuggingFace/Gradio images (update to allow UNO-FLUX output URLs)
app.get('/proxy-image', async (req, res) => {
    const { url, proxied } = req.query;
    if (!url) {
        return res.status(400).send('Missing url parameter');
    }
    // Allow proxying from yanze-pulid-flux.hf.space and UNO-FLUX output URLs
    if (!url.startsWith('https://yanze-pulid-flux.hf.space/') && !url.startsWith('https://hf.space/')) {
        return res.status(403).send('Forbidden: Only allowed to proxy from yanze-pulid-flux.hf.space or hf.space');
    }
    try {
        const response = await fetch(url, {
            method: 'GET',
            headers: {
                'Accept': 'image/avif,image/webp,image/apng,image/svg+xml,image/*,*/*;q=0.8',
                'Accept-Language': 'en-GB,en-US;q=0.9,en;q=0.8,hi;q=0.7',
                'Cache-Control': 'no-cache',
                'Pragma': 'no-cache',
                'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/138.0.0.0 Safari/537.36',
            }
        });
        if (!response.ok) {
            return res.status(response.status).send('Failed to fetch image');
        }
        res.set('Access-Control-Allow-Origin', '*');
        res.set('Access-Control-Allow-Credentials', 'true');
        res.set('Content-Type', response.headers.get('content-type') || 'image/png');
        res.set('Content-Disposition', response.headers.get('content-disposition') || 'inline');
        // If proxied param is set, respond with the proxied URL
        if (proxied === '1') {
            const proxiedUrl = `https://shashwatidr-vision.hf.space/proxy-image?url=${encodeURIComponent(url)}`;
            return res.json({ proxiedUrl });
        }
        response.body.pipe(res);
    } catch (err) {
        res.status(500).send('Proxy error: ' + err.message);
    }
});

// UNO-FLUX model generation endpoint
app.post('/api/generate/uno-flux', async (req, res) => {
    try {
        const {
            prompt = "handsome woman in the city",
            width = 512,
            height = 512,
            guidance = 4,
            num_steps = 25,
            seed = -1,
            image_prompt1,
            image_prompt2,
            image_prompt3,
            image_prompt4
        } = req.body;

        // Helper to fetch image as blob
        async function fetchImageBlob(url) {
            const response = await fetch(url);
            if (!response.ok) throw new Error('Failed to fetch image: ' + url);
            return await response.blob();
        }

        // Fetch all image prompts as blobs
        const [img1, img2, img3, img4] = await Promise.all([
            fetchImageBlob(image_prompt1),
            fetchImageBlob(image_prompt2),
            fetchImageBlob(image_prompt3),
            fetchImageBlob(image_prompt4)
        ]);

        // Connect to UNO-FLUX Gradio client
        const client = await Client.connect("bytedance-research/UNO-FLUX");
        const result = await client.predict("/gradio_generate", {
            prompt,
            width,
            height,
            guidance,
            num_steps,
            seed,
            image_prompt1: img1,
            image_prompt2: img2,
            image_prompt3: img3,
            image_prompt4: img4
        });

        res.json({ success: true, data: result.data });
    } catch (err) {
        res.status(500).json({ success: false, error: err.message });
    }
});

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