'use strict'

const dotenv = require('dotenv');
const fs = require('fs').promises;

const HfInference = require('@huggingface/inference').HfInference;

dotenv.config();

const inference = new HfInference(process.env.HF_TOKEN);

const REPO_NAME = "black-forest-labs/FLUX.1-schnell"
const IMAGE_SIZES = {
  "square": {
    height: 512,
    width: 512
  },
  "portrait-3_4": {
    height: 512,
    width: 384
  },
  "portrait-9_16": {
    height: 512,
    width: 288
  },
  "landscape-4_3": {
    height: 384,
    width: 512
  },
  "landscape-16_9": {
    height: 288,
    width: 512
  }
}

module.exports = async function (fastify, opts) {
  fastify.get('/:inputs', async function (request, reply) {
    let { inputs } = request.params;
    const { format } = request.query;
    if (format) {
      inputs = inputs + " " + format;
    }

    const slug = inputs.replace(/[^a-zA-Z0-9-_ ]/g, "").replace(/ /g, "-");

    const file = await fs.readFile(process.env.PUBLIC_FILE_UPLOAD_DIR + "/" + slug + ".png")?.catch(() => null)
    if (file) {
      return reply
        .header('Content-Type', 'image/jpeg')
        .send(file);
    }

    const { height, width } = IMAGE_SIZES[format ?? "square"] ?? IMAGE_SIZES["square"];

    const hfRequest = await inference.textToImage({
      inputs,
      model: REPO_NAME,
      parameters: {
        height,
        width
      }
    })

    const buffer = await hfRequest.arrayBuffer();
    const array = new Uint8Array(buffer);

    const dir = await fs.opendir(process.env.PUBLIC_FILE_UPLOAD_DIR).catch(() => null)
    if (!dir) await fs.mkdir(process.env.PUBLIC_FILE_UPLOAD_DIR)
    await fs.writeFile(process.env.PUBLIC_FILE_UPLOAD_DIR + "/" + slug + ".png", array)

    return reply
      .header('Content-Type', 'image/jpeg')
      .send(array);
  })
}