import { nets } from "@vladmandic/face-api";
import * as faceapi from "@vladmandic/face-api";
import canvas from "canvas";
import fs from "fs";
import * as path from "path";
import dotenv from "dotenv-extended";
import { getFaceDetectorOptions } from "../src/common";
import * as mime from "mime-types";
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 modelDir = process.env.TRAINED_MODEL_DIR as string;
  const faceDetectionNet = nets.ssdMobilenetv1;

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

  const raw = fs.readFileSync(path.join(modelDir, "data.json"), "utf-8");
  const content = JSON.parse(raw);
  const matcher = faceapi.FaceMatcher.fromJSON(content);

  let dir = path.join(process.env.OUT_DIR as string);
  const files = fs.readdirSync(dir);
  await Promise.all(
    files.map(async (file: string) => {
      const mimeType = mime.contentType(path.extname(path.join(dir, file)));
      if (!mimeType || !mimeType.startsWith("image")) {
        return;
      }
      const imgData = fs.readFileSync(path.join(dir, file));
      const input = ((await canvas.loadImage(imgData)) as unknown) as ImageData;
      const out = faceapi.createCanvasFromMedia(input);
      const resultsQuery = await faceapi
        .detectAllFaces(out, getFaceDetectorOptions(faceDetectionNet))
        .withFaceLandmarks()
        .withFaceDescriptors();

      if (resultsQuery.length > 0) {
        for (const res of resultsQuery) {
          const bestMatch = matcher.findBestMatch(res.descriptor);
          console.log(
            `Face Detected with ${
              res.detection.score * 100
            }% accuracy and a distance of ${bestMatch.distance}: ${
              bestMatch.label
            } in file ${path.join(dir, file)}`
          );
        }
      } else {
        console.log(
          `No faces detected in file ${path.join(
            process.env.REF_IMAGE_DIR as string,
            "aline",
            file
          )}`
        );
      }
    })
  );
};

main();