const axios = require('axios');
const fs = require('fs');
const path = require('path');

async function fluxImage(prompt, width = 1024, height = 1024, server = "Azure Lite Supercomputer Server") {
  try {
    const { data: init } = await axios.post(
      "https://nihalgazi-flux-unlimited.hf.space/gradio_api/call/generate_image",
      { data: [prompt, width, height, 3, true, server] },
      {
        headers: {
          "Content-Type": "application/json",
          Origin: "https://chrunos.com",
          Referer: "https://chrunos.com/",
        },
      }
    );

    const eventId = init.event_id;
    if (!eventId) throw new Error("Failed to obtain event_id.");

    const streamUrl = `https://nihalgazi-flux-unlimited.hf.space/gradio_api/call/generate_image/${eventId}`;
    let imageUrl = null;

    for (let i = 0; i < 15; i++) {
      const { data } = await axios.get(streamUrl, {
        headers: { Accept: "text/event-stream" },
      });

      const match = data.match(/"url":\s*"([^"]+)"/);
      if (match) {
        imageUrl = match[1];
        break;
      }

      await new Promise(r => setTimeout(r, 2000));
    }

    if (!imageUrl) throw new Error("Failed to retrieve image URL from stream.");
    return imageUrl;
  } catch (err) {
    throw new Error(`fluxImage error: ${err.message}`);
  }
}

async function downloadAndSaveImage(imageUrl, fileId) {
  const tmpDir = path.join(process.cwd(), 'tmp');
  
  if (!fs.existsSync(tmpDir)) {
    fs.mkdirSync(tmpDir, { recursive: true });
  }

  const filePath = path.join(tmpDir, `flux-${fileId}.webp`);

  const response = await axios.get(imageUrl, {
    responseType: 'arraybuffer',
    timeout: 30000
  });

  fs.writeFileSync(filePath, Buffer.from(response.data));

  setTimeout(() => {
    if (fs.existsSync(filePath)) {
      fs.unlinkSync(filePath);
    }
  }, 60000);

  return `flux-${fileId}.webp`;
}

const handler = async (req, res) => {
  try {
    const { prompt, key, width, height, server } = req.query;

    if (!prompt) {
      return res.status(400).json({
        author: "Herza",
        success: false,
        error: 'Missing required parameter: prompt'
      });
    }

    if (!key) {
      return res.status(400).json({
        author: "Herza",
        success: false,
        error: 'Missing required parameter: key'
      });
    }

    const imgWidth = width ? parseInt(width) : 1024;
    const imgHeight = height ? parseInt(height) : 1024;
    const selectedServer = server || "NSFW-Core: Uncensored Server 2";

    const imageUrl = await fluxImage(prompt, imgWidth, imgHeight, selectedServer);
    
    if (!imageUrl) {
      return res.status(500).json({
        author: "Herza",
        success: false,
        error: 'Failed to generate image'
      });
    }

    const fileId = Date.now() + '_' + Math.random().toString(36).substring(7);
    const filename = await downloadAndSaveImage(imageUrl, fileId);

    const baseUrl = `${req.protocol}://${req.get('host')}`;
    const localUrl = `${baseUrl}/tmp/${filename}`;

    res.json({
      author: "Herza",
      success: true,
      data: {
        result: localUrl
      }
    });

  } catch (error) {
    res.status(500).json({
      author: "Herza",
      success: false,
      error: error.message
    });
  }
};

module.exports = {
  name: 'Flux Image Generator',
  description: 'Generate images using Flux AI model',
  type: 'GET',
  routes: ['api/AI/Flux'],
  tags: ['ai', 'image', 'flux', 'generator'],
  main: ['AI'],
  parameters: ['prompt', 'key', 'width', 'height', 'server'],
  enabled: true,
  handler
};