import * as faceapi from "@vladmandic/face-api";
import canvas from "canvas";
import fs, { lstatSync } from "fs";
import * as path from "path";
import { LabeledFaceDescriptors, 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);

  const refs: Array<LabeledFaceDescriptors> = [];
  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
            .detectSingleFace(referenceImage as TNetInput, options)
            .withFaceLandmarks()
            .withFaceDescriptor();
          if (!descriptor || !descriptor.descriptor) {
            throw new Error("No face found");
          }

          const faceDescriptors = [descriptor.descriptor];
          return new faceapi.LabeledFaceDescriptors(dir, faceDescriptors);
        } catch (err) {
          console.log(
            "An error occurred loading image at path: " +
              path.join(inputDir, dir, file)
          );
        }
        return undefined;
      })
    );

    if (referenceResults) {
      refs.push(
        ...(referenceResults.filter((e) => e) as LabeledFaceDescriptors[])
      );
    }
  }

  const faceMatcher = new faceapi.FaceMatcher(refs);

  fs.writeFile(
    path.join(outDir, "data.json"),
    JSON.stringify(faceMatcher.toJSON()),
    "utf8",
    (err) => {
      if (err) {
        console.log(`An error occurred while writing data model to file`);
      }

      console.log(`Successfully wrote data model to file`);
    }
  );
};

main();