#!/usr/bin/env python

import subprocess
import threading
import getpass
import os
import sys
import signal
import time
import glob

class I3Locker:
    def __init__(self, lockpid):
        self.killed = False
        self.lockpid = lockpid

    def run(self):
        died = False
        while not died and not self.killed:
            try:
                os.kill(self.lockpid, 0)
                time.sleep(0.5)
            except:
                died = True

        if self.killed:
            return 1
        else:
            return 0

    def kill(self):
        self.killed = True
        os.kill(self.lockpid, signal.SIGTERM)

class FaceLocker:
    def __init__(self):
        self.delay = 200
        self.dev = 2
        self.running = False
        self.killed = False
        self.waitingProc = None

    def run(self):
        self.running = True

        # Import here because it's sloow
        import face_recognition
        import cv2
        import numpy as np

        # 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:
            self.waitForKey(keyboard, key)

        if self.killed:
            return 0

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

        if self.killed:
            return 1
        elif self.matching:
            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):
                    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):
        if self.killed:
            return

        print("Waiting for enter before starting face recognition")
        self.waitingProc = subprocess.Popen(
            f"xinput test '{keyboard}' | grep --line-buffered 'key press   {key}' | exit",
            shell=True)

        # Blink IR blasters
        import cv2
        cap = cv2.VideoCapture(self.dev)
        time.sleep(0.1)
        cap.release()

        self.waitingProc.wait()

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

lockers = [
    I3Locker(int(sys.argv[1])),
    FaceLocker(),
]

def runLocker(locker):
    print("Starting "+locker.__class__.__name__)
    ret = locker.run()
    if ret == 0:
        print(locker.__class__.__name__+" unlocked.")
        for l in lockers:
            if l == locker:
                continue
            l.kill()
    elif ret != 1:
        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()