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 idlivesdk import getHWID
from idlivesdk import setLicenseKey
from idlivesdk import initSDK
from idlivesdk import processImage

licenseKeyPath = "license.txt"
license = os.environ.get("LICENSE_KEY")

if license is None:
    try:
        with open(licenseKeyPath, 'r') as file:
            license = file.read().strip()
    except IOError as exc:
        print("failed to open license.txt: ", exc.errno)
    print("License Key: ", license)

hwid = getHWID()
print("HWID: ", hwid.decode('utf-8'))

ret = setLicenseKey(license.encode('utf-8'))
print("Set License: ", ret)

ret = initSDK("model".encode('utf-8'))
print("Init: ", ret)

app = Flask(__name__) 

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("No EXIF data found")

    return image


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

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

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

    image_np = np.asarray(image)
    result = processImage(image_np, image_np.shape[1], image_np.shape[0])

    if result is None:
        result = "Failed to process image"
        response = jsonify({"resultCode": "Error", "result": result})

        response.status_code = 200
        response.headers["Content-Type"] = "application/json; charset=utf-8"
        return response
    else:
        result_dict = json.loads(result.decode('utf-8'))
        response = jsonify({"resultCode": "Ok", "result": result_dict})

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

@app.route('/process_image_base64', methods=['POST'])
def process_image_base64():
    try:
        content = request.get_json()
        base64_image = content['base64']

        image_data = base64.b64decode(base64_image)
        image = apply_exif_rotation(Image.open(io.BytesIO(image_data))).convert("RGB")
    except:
        result = "Failed to parse base64"
        response = jsonify({"resultCode": "Error", "result": result})

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

    image_np = np.asarray(image)
    result = processImage(image_np, image_np.shape[1], image_np.shape[0])

    if result is None:
        result = "Failed to process image"
        response = jsonify({"resultCode": "Error", "result": result})

        response.status_code = 200
        response.headers["Content-Type"] = "application/json; charset=utf-8"
        return response
    else:
        result_dict = json.loads(result.decode('utf-8'))
        response = jsonify({"resultCode": "Ok", "result": result_dict})

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

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