/**
 * Z-Image-Turbo API Wrapper - Node.js Express Version
 * 
 * This server wraps the Gradio API and provides a simple REST endpoint
 * that returns the image URL directly
 */

const express = require('express');
const cors = require('cors');
const axios = require('axios');
const bodyParser = require('body-parser');

const app = express();
const PORT = process.env.PORT || 5000;
const GRADIO_API_URL = 'https://mohamedislegend4-z-image-turbo-api.hf.space';
const MAX_POLL_ATTEMPTS = 120;
const POLL_INTERVAL = 1000; // ms

// Middleware
app.use(cors());
app.use(bodyParser.json());

// Logger
const log = (level, message, data = '') => {
  const timestamp = new Date().toISOString();
  console.log(`[${timestamp}] ${level}: ${message} ${data}`);
};

/**
 * Call a Gradio endpoint
 */
async function callGradioEndpoint(endpoint, params) {
  const url = `${GRADIO_API_URL}/gradio_api/call/v2/${endpoint}`;
  
  log('INFO', `Calling endpoint: ${endpoint}`, JSON.stringify(params));
  
  try {
    const response = await axios.post(url, params, {
      headers: { 'Content-Type': 'application/json' },
      timeout: 30000
    });
    
    log('INFO', 'Response received', JSON.stringify(response.data));
    return response.data;
  } catch (error) {
    log('ERROR', `Failed to call endpoint: ${error.message}`);
    throw error;
  }
}

/**
 * Poll for result
 */
async function pollResult(endpoint, eventId) {
  const url = `${GRADIO_API_URL}/gradio_api/call/${endpoint}/${eventId}`;
  
  log('INFO', `Polling result: ${endpoint}/${eventId}`);
  
  for (let attempt = 0; attempt < MAX_POLL_ATTEMPTS; attempt++) {
    try {
      const response = await axios.get(url, {
        timeout: 30000
      });
      
      const data = response.data;
      
      // Check if result is ready
      if (data.result && data.result !== null) {
        log('INFO', `Got result on attempt ${attempt + 1}`);
        return data.result;
      }
      
      log('INFO', `Attempt ${attempt + 1}: No result yet, waiting...`);
      
      // Wait before next poll
      await new Promise(resolve => setTimeout(resolve, POLL_INTERVAL));
      
    } catch (error) {
      log('WARN', `Poll request failed: ${error.message}`);
      await new Promise(resolve => setTimeout(resolve, POLL_INTERVAL));
    }
  }
  
  log('ERROR', `Polling timeout after ${MAX_POLL_ATTEMPTS} attempts`);
  return null;
}

/**
 * Generate image using Z-Image-Turbo
 */
async function generateImage(prompt, steps = 20, height = 512, width = 512) {
  try {
    const params = {
      prompt,
      steps: Math.min(Math.max(steps, 1), 50),
      height,
      width
    };
    
    // Call the endpoint
    const callResponse = await callGradioEndpoint('generate', params);
    
    const eventId = callResponse.event_id;
    if (!eventId) {
      throw new Error('No event_id in response');
    }
    
    // Poll for result
    const result = await pollResult('generate', eventId);
    
    if (!result) {
      throw new Error('Failed to get result from API');
    }
    
    // Extract image data
    if (Array.isArray(result) && result.length > 0) {
      return result[0];
    }
    
    return result;
    
  } catch (error) {
    log('ERROR', `Error generating image: ${error.message}`);
    throw error;
  }
}

/**
 * Routes
 */

// Health check
app.get('/health', (req, res) => {
  res.json({ status: 'ok' });
});

// API Info
app.get('/api/info', (req, res) => {
  res.json({
    api_name: 'Z-Image-Turbo API Wrapper',
    version: '1.0.0',
    language: 'Node.js',
    endpoints: [
      {
        path: '/api/generate',
        method: 'POST',
        description: 'Generate an image from a text prompt',
        parameters: {
          prompt: {
            type: 'string',
            required: true,
            description: 'Text description of the image'
          },
          steps: {
            type: 'integer',
            required: false,
            default: 20,
            min: 1,
            max: 50,
            description: 'Number of inference steps'
          },
          height: {
            type: 'integer',
            required: false,
            default: 512,
            description: 'Image height in pixels'
          },
          width: {
            type: 'integer',
            required: false,
            default: 512,
            description: 'Image width in pixels'
          }
        }
      }
    ]
  });
});

// Generate image endpoint
app.post('/api/generate', async (req, res) => {
  try {
    const { prompt, steps = 20, height = 512, width = 512 } = req.body;
    
    // Validate input
    if (!prompt) {
      return res.status(400).json({
        success: false,
        error: 'Missing required parameter: prompt'
      });
    }
    
    log('INFO', `Generating image for: ${prompt}`);
    
    // Generate image
    const imageData = await generateImage(prompt, steps, height, width);
    
    if (!imageData) {
      return res.status(500).json({
        success: false,
        error: 'Failed to generate image'
      });
    }
    
    // Build response
    const result = {
      success: true,
      prompt,
      steps: Math.min(Math.max(steps, 1), 50),
      height,
      width
    };
    
    // Add image data fields
    if (typeof imageData === 'object') {
      if (imageData.url) result.image_url = imageData.url;
      if (imageData.path) result.image_path = imageData.path;
      if (imageData.size) result.size = imageData.size;
      if (imageData.mime_type) result.mime_type = imageData.mime_type;
      if (imageData.orig_name) result.filename = imageData.orig_name;
    }
    
    res.json(result);
    
  } catch (error) {
    log('ERROR', `Error in /api/generate: ${error.message}`);
    res.status(500).json({
      success: false,
      error: error.message
    });
  }
});

// 404 handler
app.use((req, res) => {
  res.status(404).json({ error: 'Not found' });
});

// Error handler
app.use((err, req, res, next) => {
  log('ERROR', `Unhandled error: ${err.message}`);
  res.status(500).json({
    error: 'Internal server error',
    message: process.env.NODE_ENV === 'development' ? err.message : undefined
  });
});

// Start server
app.listen(PORT, '0.0.0.0', () => {
  log('INFO', '='.repeat(60));
  log('INFO', 'Z-Image-Turbo API Wrapper (Node.js)');
  log('INFO', '='.repeat(60));
  log('INFO', `Gradio API URL: ${GRADIO_API_URL}`);
  log('INFO', `Server running on: http://localhost:${PORT}`);
  log('INFO', `Health check: http://localhost:${PORT}/health`);
  log('INFO', `API info: http://localhost:${PORT}/api/info`);
  log('INFO', `Generate endpoint: POST http://localhost:${PORT}/api/generate`);
  log('INFO', '='.repeat(60));
});

module.exports = app;