import sys
sys.path.append('.')

import os
import numpy as np
import base64
import json
import io
from PIL import Image, ExifTags
from flask import Flask, request, jsonify
from engine.header import *

maxFaceCount1 = 5
maxFaceCount2 = 10

file_path = os.path.abspath(__file__)
root_path = os.path.dirname(file_path)

app = Flask(__name__) 
app.config['SITE'] = "http://0.0.0.0:8000/"
app.config['DEBUG'] = False

device_id = getHWID().decode('utf-8')
print_info('\t <Hardware ID> \t\t {}'.format(device_id))

license = os.environ.get("FR_LICENSE_KEY")
dict_path = os.path.join(root_path, "engine")

ret = -1
if license is None:
    try:
        licenseKeyPath = os.path.join(root_path, "license.txt")
        with open(licenseKeyPath, 'r') as file:
            license = file.read().strip()
    except IOError as exc:
        print_error(f"failed to open license.txt: {exc.errno}")

print_log(f"License Key: \n{license}")
ret = setLicenseKey(license.encode('utf-8'))
print_log(f"Set License: {ret}")

ret = initSDK(dict_path.encode('utf-8'))
if ret == 0:
    print_log("Successfully init SDK!")
else:
    print_error(f"Failed to init SDK, Error code {ret}")
    exit(-1)

def apply_exif_rotation(image):
    try:
        exif = image._getexif()
        if exif is not None:
            for orientation in ExifTags.TAGS.keys():
                if ExifTags.TAGS[orientation] == 'Orientation':
                    break

            # Get the orientation value
            orientation = exif.get(orientation, None)

            # Apply the appropriate rotation based on the orientation
            if orientation == 3:
                image = image.rotate(180, expand=True)
            elif orientation == 6:
                image = image.rotate(270, expand=True)
            elif orientation == 8:
                image = image.rotate(90, expand=True)

    except AttributeError:
        print_error("No EXIF data found")

    return image

@app.route('/compare_face', methods=['POST'])
def compare_face():
    file1 = request.files['image1']
    file2 = request.files['image2']

    try:
        image1 = apply_exif_rotation(Image.open(file1)).convert('RGB')        
    except:
        result = "Failed to open file1"
        response = jsonify({"resultCode": result})

        response.status_code = 200
        response.headers["Content-Type"] = "application/json; charset=utf-8"
        return response

    try:
        image2 = apply_exif_rotation(Image.open(file2)).convert('RGB')        
    except:
        result = "Failed to open file2"
        response = jsonify({"resultCode": result})

        response.status_code = 200
        response.headers["Content-Type"] = "application/json; charset=utf-8"
        return response

    image_np1 = np.asarray(image1)
    faces1 = (Face * maxFaceCount1)()
    faceNum1 = processImage(image_np1, image_np1.shape[1], image_np1.shape[0], faces1, maxFaceCount1)

    faces1_result = []
    for i in range(faceNum1):
        face = {"x": faces1[i].x, "y": faces1[i].y, "width": faces1[i].width, "height": faces1[i].height}
        faces1_result.append(face)

    image_np2 = np.asarray(image2)
    faces2 = (Face * maxFaceCount1)()
    faceNum2 = processImage(image_np2, image_np2.shape[1], image_np2.shape[0], faces2, maxFaceCount1)

    faces2_result = []
    for i in range(faceNum2):
        face = {"x": faces2[i].x, "y": faces2[i].y, "width": faces2[i].width, "height": faces2[i].height}
        faces2_result.append(face)

    if faceNum1 > 0 and faceNum2 > 0:
        results = []
        for i in range(faceNum1):
            for j in range(faceNum2):                
                score = verifyFeat(faces1[i].featSize, faces1[i].featData, faces2[j].featSize, faces2[j].featData)
                match_result = {"face1": i, "face2": j, "score": score}
                results.append(match_result)

        response = jsonify({"resultCode": "Ok", "faces1": faces1_result, "faces2": faces2_result, "results": results})

        response.status_code = 200
        response.headers["Content-Type"] = "application/json; charset=utf-8"
        return response
    elif faceNum1 == 0:
        response = jsonify({"resultCode": "No face1", "faces1": faces1, "faces2": faces2})

        response.status_code = 200
        response.headers["Content-Type"] = "application/json; charset=utf-8"
        return response
    elif faceNum2 == 0:
        response = jsonify({"resultCode": "No face2", "faces1": faces1, "faces2": faces2})

        response.status_code = 200
        response.headers["Content-Type"] = "application/json; charset=utf-8"
        return response

@app.route('/compare_face_base64', methods=['POST'])
def compare_face_base64():
    base64_image1 = ""
    base64_image2 = ""

    try:
        content = request.get_json()
        base64_image1 = content['image1_base64']

        image_data1 = base64.b64decode(base64_image1)
        image1 = apply_exif_rotation(Image.open(io.BytesIO(image_data1))).convert("RGB")
    except:
        result = "Failed to parse image1 base64"
        response = jsonify({"resultCode": result})

        response.status_code = 200
        response.headers["Content-Type"] = "application/json; charset=utf-8"
        return response

    try:
        content = request.get_json()
        base64_image2 = content['image2_base64']

        image_data2 = base64.b64decode(base64_image2)
        image2 = apply_exif_rotation(Image.open(io.BytesIO(image_data2))).convert("RGB")
    except:
        result = "Failed to parse image2 base64"
        response = jsonify({"resultCode": result})

        response.status_code = 200
        response.headers["Content-Type"] = "application/json; charset=utf-8"
        return response

    image_np1 = np.asarray(image1)
    faces1 = (Face * maxFaceCount1)()
    faceNum1 = processImage(image_np1, image_np1.shape[1], image_np1.shape[0], faces1, maxFaceCount1)

    faces1_result = []
    for i in range(faceNum1):
        face = {"x": faces1[i].x, "y": faces1[i].y, "width": faces1[i].width, "height": faces1[i].height}
        faces1_result.append(face)

    image_np2 = np.asarray(image2)
    faces2 = (Face * maxFaceCount1)()
    faceNum2 = processImage(image_np2, image_np2.shape[1], image_np2.shape[0], faces2, maxFaceCount1)

    faces2_result = []
    for i in range(faceNum2):
        face = {"x": faces2[i].x, "y": faces2[i].y, "width": faces2[i].width, "height": faces2[i].height}
        faces2_result.append(face)

    if faceNum1 > 0 and faceNum2 > 0:
        results = []
        for i in range(faceNum1):
            for j in range(faceNum2):                
                score = verifyFeat(faces1[i].featSize, faces1[i].featData, faces2[j].featSize, faces2[j].featData)
                match_result = {"face1": i, "face2": j, "score": score}
                results.append(match_result)

        response = jsonify({"resultCode": "Ok", "faces1": faces1_result, "faces2": faces2_result, "results": results})

        response.status_code = 200
        response.headers["Content-Type"] = "application/json; charset=utf-8"
        return response
    elif faceNum1 == 0:
        response = jsonify({"resultCode": "No face1", "faces1": faces1, "faces2": faces2})

        response.status_code = 200
        response.headers["Content-Type"] = "application/json; charset=utf-8"
        return response
    elif faceNum2 == 0:
        response = jsonify({"resultCode": "No face2", "faces1": faces1, "faces2": faces2})

        response.status_code = 200
        response.headers["Content-Type"] = "application/json; charset=utf-8"
        return response

@app.route('/detect_face', methods=['POST'])
def detect_face():
    file = request.files['image']

    try:
        image = apply_exif_rotation(Image.open(file)).convert('RGB')
    except:
        result = "Failed to open file"
        response = jsonify({"resultCode": result})

        response.status_code = 200
        response.headers["Content-Type"] = "application/json; charset=utf-8"
        return response

    image_np = np.asarray(image)
    faces = (Face * maxFaceCount2)()
    faceNum = processImage(image_np, image_np.shape[1], image_np.shape[0], faces, maxFaceCount2)

    if faceNum > 0:
        faces_result = []
        for i in range(faceNum):
            face = {"rect": {"x": faces[i].x, "y": faces[i].y, "width": faces[i].width, "height": faces[i].height},
            "angles": {"yaw": faces[i].yaw, "roll": faces[i].roll, "pitch": faces[i].pitch},
            "age_gender": {"age": faces[i].age, "gender": faces[i].gender},
            "emotion": {"neutral": faces[i].neutral, "happy": faces[i].happy, "angry": faces[i].angry,
            "surprised": faces[i].surprised, "disgusted": faces[i].disgusted, "sad": faces[i].sad, "scared": faces[i].scared},
            "attribute": {"masked": faces[i].masked, "left_eye_opened": faces[i].left_eye_opened, "right_eye_opened": faces[i].right_eye_opened}}
            faces_result.append(face)

        response = jsonify({"resultCode": "Ok", "result": faces_result})
    else:
        response = jsonify({"resultCode": "No face"})

    response.status_code = 200
    response.headers["Content-Type"] = "application/json; charset=utf-8"
    return response

if __name__ == '__main__':
    port = int(os.environ.get("PORT", 8000))
    app.run(host='0.0.0.0', port=port)