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

dotenv.config();

const { Pool } = pg;

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

// Middleware
app.use(cors());
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 (now GET, returns only image URL)
app.get('/api/generate/flux', authenticateToken, async (req, res) => {
    try {
        const { prompt, aspect } = req.query;
        if (!prompt) {
            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 generation for prompt: "${prompt}" with session: ${sessionHash}`);

        // Join the queue using the correct API structure
        const joinResponse = await axios.post('https://yanze-pulid-flux.hf.space/queue/join', {
            data: [
                prompt,
                null, // id_image (null for no ID image)
                0,    // start_step
                4,    // guidance
                "-1", // seed
                1,    // true_cfg
                width,
                height,
                20,   // num_steps
                1,    // id_weight
                "bad quality, worst quality, text, signature, watermark, extra limbs", // neg_prompt
                1,    // timestep_to_start_cfg
                128   // max_sequence_length
            ],
            event_data: null,
            fn_index: 4,
            trigger_id: 21,
            session_hash: sessionHash
        }, {
            headers: {
                'Content-Type': 'application/json',
                'Referer': 'https://yanze-pulid-flux.hf.space/'
            }
        });

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

        // Poll for results
        let attempts = 0;
        const maxAttempts = 60;
        
        while (attempts < maxAttempts) {
            try {
                const dataResponse = await axios.get(
                    `https://yanze-pulid-flux.hf.space/queue/data?session_hash=${encodeURIComponent(sessionHash)}`,
                    {
                        headers: {
                            'Accept': 'text/event-stream',
                            'Referer': 'https://yanze-pulid-flux.hf.space/'
                        },
                        timeout: 10000
                    }
                );

                const data = dataResponse.data;
                console.log(`[Flux] Poll attempt ${attempts + 1}, data length: ${data.length}`);
                
                // Parse the event stream data
                const lines = data.split('\n');
                let eventData = '';
                
                for (const line of lines) {
                    if (line.startsWith('data: ')) {
                        eventData += line.substring(6);
                    } else if (line === '') {
                        if (eventData) {
                            try {
                                const json = JSON.parse(eventData);
                                console.log(`[Flux] Parsed event: ${json.msg}`);
                                
                                if (json.msg === 'process_completed' && json.output?.data?.[0]?.url) {
                                    const imageUrl = json.output.data[0].url;
                                    console.log(`[Flux] Found image:`, imageUrl);
                                    
                                    res.json({ success: true, imageUrl: imageUrl });
                                    return;
                                }

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

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

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

        console.log(`[Flux] Generation timed out after ${maxAttempts} attempts`);
        res.status(500).json({ success: false, error: 'Generation timed out' });

    } catch (error) {
        console.error('[Flux] Generation endpoint error:', 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' });
  }
});

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