from flask import Flask, request, redirect, url_for, flash, jsonify
from detect_object import predict_object_function
import os
from flask_wtf import CSRFProtect
from waitress import serve

app = Flask(__name__)
app.config['SECRET_KEY'] = '8BYkEfBA6O6zWlSihBXox7C0sKR6b'
csrf = CSRFProtect(app)



@app.route('/')
def home():
    return "Hello Wav2Lip - Flask API Running on Hugging Face Spaces!"


@app.route('/predict_object', methods=['POST'])
@csrf.exempt
def predict_object():
    # Receive image file from frontend
    image_file = request.files['media']

    image_path = os.path.join('predict_image', image_file.filename)
    image_file.save(image_path)

    try:
        result = predict_object_function(image_path)
        print(result)
    except RuntimeError as e:
        text = f"{str(e)})", 'danger'
        return {"response": text}
    except Exception as e:
        print(str(e))
        text = "An error occurred during prediction. Please try again"
        return {"response": text}
    finally:
        os.remove(image_path)

    imgs = result
    no_p = len(imgs)
    if no_p < 1:
        response = {"response": "No object detected."}
        return jsonify(response)
    percentage = result[1]
    # response = {'response': f"Detected {result[0]} with confidence of {round(percentage, 2)}%"}
    response = {'response': f"{result[0]}", "confidence": round(percentage, 2)}

    return jsonify(response)


if __name__ == '__main__':
    serve(app, host='0.0.0.0', port=7860)