from flask import Flask, request, jsonify, send_from_directory
from nsfw_detector import predict
import requests
import io
import random
import os
from PIL import Image
from datetime import datetime
import string
import re

app = Flask(__name__)
model = predict.load_model('./nsfw_mobilenet2.224x224.h5')

@app.route("/imagebuffer", methods=["POST"])
def check():
    try:
        data = request.get_json()
        imgdata = data["data"]

        #imgdata is like 'data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQABAAD/2wBD...', save it to a file in ./TEMP
        imgdata = re.sub('^data:image/.+;base64,', '', imgdata)
        imgdata = bytes(imgdata, encoding="utf-8")
        img = Image.open(io.BytesIO(imgdata))
        filename = ''.join(random.choices(string.ascii_uppercase + string.digits, k=10)) + ".jpg"
        img.save(f"./TEMP/{filename}")

        #predict the image
        result = predict.predict(model, f"./TEMP/{filename}")

        response = {
            "success": True,
            "result": result
        }

    except Exception as e:
        response = {
            "success": False,
            "error": str(e)
        }

    return jsonify(response)

@app.route("/imageurl", methods=["POST"])
def checkUrl():
    try:
        data = request.get_json()
        url = data["url"]

        #download the image
        response = requests.get(url)
        img = Image.open(io.BytesIO(response.content))
        filename = ''.join(random.choices(string.ascii_uppercase + string.digits, k=10)) + ".jpg"
        img.save(f"./TEMP/{filename}")

        #predict the image
        result = predict.predict(model, f"./TEMP/{filename}")

        response = {
            "success": True,
            "result": result
        }

    except Exception as e:
        response = {
            "success": False,
            "error": str(e)
        }

    return jsonify(response)


if __name__ == "__main__":
    if not os.path.exists("./TEMP"):
        os.makedirs("./TEMP")

    app.run(debug=True, host="0.0.0.0", port=7860)