import * as faceapi from "@vladmandic/face-api";
import canvas from "canvas";
import fs, { lstatSync } from "fs";
import * as path from "path";
import { TNetInput } from "@vladmandic/face-api";
import * as mime from "mime-types";
import dotenv from "dotenv-extended";
import { getFaceDetectorOptions } from "../src/common";
require("@tensorflow/tfjs-node");

const { Canvas, Image, ImageData } = canvas;
//@ts-ignore
faceapi.env.monkeyPatch({ Canvas, Image, ImageData });

const main = async () => {
  dotenv.load({
    silent: false,
    errorOnMissing: true,
  });
  const inputDir = process.env.REF_IMAGE_DIR as string;
  const outDir = process.env.TRAINED_MODEL_DIR as string;

  const faceDetectionNet = faceapi.nets.ssdMobilenetv1;
  await faceDetectionNet.loadFromDisk(path.join(__dirname, "../weights"));
  await faceapi.nets.faceLandmark68Net.loadFromDisk(
    path.join(__dirname, "../weights")
  );
  await faceapi.nets.faceRecognitionNet.loadFromDisk(
    path.join(__dirname, "../weights")
  );

  const options = getFaceDetectorOptions(faceDetectionNet);

  const dirs = fs.readdirSync(inputDir);

  for (const dir of dirs) {
    if (!lstatSync(path.join(inputDir, dir)).isDirectory()) {
      continue;
    }
    const files = fs.readdirSync(path.join(inputDir, dir));
    let referenceResults = await Promise.all(
      files.map(async (file: string) => {
        const mimeType = mime.contentType(
          path.extname(path.join(inputDir, dir, file))
        );
        if (!mimeType || !mimeType.startsWith("image")) {
          return;
        }
        console.log(path.join(inputDir, dir, file));

        try {
          const referenceImage = (await canvas.loadImage(
            path.join(inputDir, dir, file)
          )) as unknown;

          const descriptor = await faceapi
            .detectAllFaces(referenceImage as TNetInput, options)
            .withFaceLandmarks()
            .withFaceDescriptors();

          return descriptor.length > 0 ? descriptor : undefined;
        } catch (err) {
          console.log(
            "An error occurred loading image at path: " +
              path.join(inputDir, dir, file)
          );
        }
        return undefined;
      })
    );

    const items = [];
    for (const item of referenceResults) {
      if (item) {
        items.push(...item);
      }
    }
    const faceMatcher = new faceapi.FaceMatcher(items);
    fs.writeFile(
      path.join(outDir, dir + ".json"),
      JSON.stringify(faceMatcher.toJSON()),
      "utf8",
      (err) => {
        if (err) {
          console.log(`An error occurred while writing ${dir} model to file`);
        }

        console.log(`Successfully wrote ${dir} model to file`);
      }
    );
  }
};

main();