"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: {
    width: 1024,
    height: 1024,
  },
  "portrait-3_4": {
    width: 768,
    height: 1024,
  },
  "portrait-9_16": {
    width: 576,
    height: 1024,
  },
  "landscape-4_3": {
    width: 1024,
    height: 768,
  },
  "landscape-16_9": {
    width: 1024,
    height: 576,
  },
};

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);
  });
};