import {
  SlashCommandBuilder,
  AttachmentBuilder,
  ChatInputCommandInteraction,
} from 'discord.js';
import axios from 'axios';
import Replicate, { Prediction } from 'replicate';
import config from '../config';

const replicate = new Replicate({
  auth: config.replicateApiKey,
});

// Initialise the command data.
export const data = new SlashCommandBuilder()
  .setName('image')
  .setDescription('Generate an image based on a prompt.')
  .addStringOption((option) =>
    option
      .setName('prompt')
      .setDescription('The prompt to generate an image from')
      .setRequired(true)
  );

console.log(`Loaded ${data.name} command.`);

/**
 * Generate an image based on a prompt and send it back to the user.
 * @param interaction The interaction that triggered the command.
 * @returns A promise that resolves when the command is finished executing.
 */
export async function execute(
  interaction: ChatInputCommandInteraction
): Promise<void> {
  await interaction.deferReply();

  const prompt = interaction.options.get('prompt')?.value as string;

  try {
    // Create image generation prediction
    const prediction = await replicate.predictions.create({
      model: 'black-forest-labs/flux-schnell',
      input: { prompt },
    });

    // Poll until the image generation is complete
    const completedPrediction = await pollPredictionStatus(prediction.id);

    if (!completedPrediction || !completedPrediction.output) {
      throw new Error('Failed to generate the image.');
    }

    const imageUrl = completedPrediction.output[0];

    // Download the generated image
    const imageBuffer = await downloadImage(imageUrl);

    // Create an attachment to send the image back to the user
    const attachment = new AttachmentBuilder(imageBuffer, {
      name: 'image.png',
    });

    // Edit the deferred reply to include the generated image
    await interaction.editReply({ files: [attachment] });
  } catch (error: any) {
    // Provide a more informative error message to the user
    console.error(error); // Log the error for debugging purposes
    await interaction.editReply(`An error occurred: ${error.message}`);
  }
}

/**
 * Poll the status of a prediction until it is no longer in progress.
 * @param predictionId The ID of the prediction to poll.
 * @param maxAttempts The maximum number of attempts to poll the prediction.
 * @param interval The interval between each polling attempt in milliseconds.
 * @returns The final status of the prediction.
 */
async function pollPredictionStatus(
  predictionId: string,
  maxAttempts = 5,
  interval = 2000
): Promise<Prediction> {
  for (let attempt = 0; attempt < maxAttempts; attempt++) {
    const latestPrediction = await replicate.predictions.get(predictionId);

    if (
      latestPrediction.status !== 'starting' &&
      latestPrediction.status !== 'processing'
    ) {
      return latestPrediction;
    }

    // Wait before checking again
    await new Promise((resolve) => setTimeout(resolve, interval));
  }

  throw new Error('Prediction timed out.');
}

/**
 * Download an image from a URL and return it as a Buffer.
 * @param url The URL of the image to download.
 * @returns A Buffer containing the downloaded image.
 */
async function downloadImage(url: string): Promise<Buffer> {
  try {
    const response = await axios.get(url, { responseType: 'arraybuffer' });
    return Buffer.from(response.data);
  } catch (error) {
    throw new Error('Failed to download the image.');
  }
}