#!/usr/bin/env python

import subprocess
import threading
import getpass
import os
import time
import glob
import cv2
import numpy as np

class I3Locker:
    def run(self):
        self.proc = subprocess.Popen([ "mlock" ])
        code = self.proc.wait()
        if code == 0 or self.killed:
            return 0
        else:
            print("mlock exited with code "+str(code))
            return -1

    def kill(self):
        self.killed = True
        self.proc.terminate()

class FaceLocker:
    def run(self):
        self.delay = 200
        self.dev = 2
        self.running = True
        self.waitingProc = None

        # Import here because it's sloow
        import face_recognition

        # Read all face files
        faceencs = []
        path = f"./faces/{getpass.getuser()}/*.npy"
        paths = []
        for p in glob.glob(path):
            print(f"reading {p}")
            faceencs.append(np.load(p))
            paths.append(p)

        # Wait here if we're on battery
        battery = "/sys/class/power_supply/BAT0"
        keyboard = "AT Translated Set 2 keyboard"
        key = 36
        bat = False
        with open(f"{battery}/status", "r") as f:
            s = f.read().strip()
            if s == "Discharging" or s == "Unknown":
                bat = True
        if bat:
            print("Waiting for enter before starting face recognition")
            self.waitForKey(keyboard, key)

        # Match faces, blocks until a match is found or we're killed
        self.runFaces(faceencs, paths, np, face_recognition, cv2)

        if self.matching or self.killed:
            return 0
        else:
            return -1

    def runFaces(self, faceencs, paths, np, face_recognition, cv2):
        self.matching = False
        cap = cv2.VideoCapture(self.dev)
        tacc = self.delay
        then = 0
        avg = 128
        while not self.matching and self.running:
            ret, frame = cap.read()
            mean = cv2.mean(frame)[0]
            avg = (avg + mean) / 2
            if mean < avg:
                continue

            # delay
            now = time.time() * 1000
            if tacc < self.delay:
                tacc += now - then
                then = now
                continue
            else:
                tacc = 0
                then = now

            scale = 1
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
            small_rgb_frame = cv2.resize(rgb_frame, (0, 0), fx=scale, fy=scale)

            framelocs = face_recognition.face_locations(small_rgb_frame)
            frameencs = face_recognition.face_encodings(small_rgb_frame, framelocs)

            # Loop through each face in this frame of video
            for (top, right, bottom, left), frameenc in zip(framelocs, frameencs):
                # See if the face is a match for the known face(s)
                dists = face_recognition.face_distance(faceencs, frameenc)
                dist = dists[0]
                distidx = 0
                for i, d in enumerate(dists):
                    print(i, d)
                    if d < dist:
                        dist = d
                        distidx = i
                print(f"Distance: {dist} ({paths[distidx]})")

                # If a match was found in known_face_encodings, just use the first one.
                if dist <= 0.4:
                    self.matching = True

    def waitForKey(self, keyboard, key):
        self.waitingProc = subprocess.Popen(
            f"xinput test '{keyboard}' | grep --line-buffered 'key press   {key}' | exit",
            shell=True)

        # Blink IR blasters
        cap = cv2.VideoCapture(self.dev)
        cap.release()

        self.waitingProc.wait()

    def kill(self):
        self.killed = True
        self.running = False
        if self.waitingProc:
            self.waitingProc.terminate()

lockers = [
    I3Locker(),
    FaceLocker(),
]

def runLocker(locker):
    ret = locker.run()
    if ret == 0:
        print(locker.__class__.__name__+" unlocked.")
        for l in lockers:
            if l == locker:
                continue
            l.kill()
    else:
        print(locker.__class__.__name__+" failed.")


threads = []
for locker in lockers:
    th = threading.Thread(target=runLocker, args=(locker,))
    th.start()
    threads.append(th)

for th in threads:
    th.join()