import os
# Set a custom directory for InsightFace models
os.environ['INSIGHTFACE_HOME'] = 'models'
os.makedirs('models', exist_ok=True)

from insightface.app import FaceAnalysis
import insightface

# Now initialize your model
app_model = FaceAnalysis(name='buffalo_l') # Todo: Evaluate how buffalo_m performs
try:
    app_model.prepare(ctx_id=0, det_size=(640, 640))  # Try GPU
except:
    app_model.prepare(ctx_id=-1, det_size=(640, 640))  # Fallback to CPU



from flask import Flask, render_template, Response, request, jsonify, send_file, abort
import cv2
import numpy as np
import base64
from insightface.app import FaceAnalysis
import os
from werkzeug.utils import secure_filename
import shutil

app = Flask(__name__)

# Get the access key from Hugging Face environment variable
ACCESS_KEY = os.environ.get('KEY', 'default_access_key')

def check_auth():
    """Check if the request is authenticated with the correct key"""
    provided_key = request.args.get('key')
    if provided_key != ACCESS_KEY:
        return False
    return True

def require_auth(f):
    """Decorator to require authentication for routes"""
    def decorated_function(*args, **kwargs):
        if not check_auth():
            abort(401)  # Unauthorized
        return f(*args, **kwargs)
    decorated_function.__name__ = f.__name__
    return decorated_function

# Initialize the face detection and swapping models
app_model = FaceAnalysis(name='buffalo_l')
app_model.prepare(ctx_id=0, det_size=(640, 640))

swapper_model_path = 'src/inswapper_128.onnx'
swapper = insightface.model_zoo.get_model(swapper_model_path, download=False, download_zip=False)

# Load the default image for face swapping
default_img_path = 'assets/face.jpg'
default_img = cv2.imread(default_img_path)

# Make sure the assets directory exists
if not os.path.exists('assets'):
    os.makedirs('assets')

# Check if a face is detected in the default image
faces_in_default_img = app_model.get(default_img)
if len(faces_in_default_img) == 0:
    raise Exception("No face detected in the default image.")

default_face = faces_in_default_img[0]


@app.route('/')
@require_auth
def index():
    return render_template('index.html')


@app.route('/process_frame', methods=['POST'])
@require_auth
def process_frame():
    try:
        data = request.json['image']
        img_data = base64.b64decode(data.split(',')[1])
        np_img = np.frombuffer(img_data, np.uint8)
        frame = cv2.imdecode(np_img, cv2.IMREAD_COLOR)

        faces_in_frame = app_model.get(frame)
        if len(faces_in_frame) > 0:
            face_in_frame = faces_in_frame[0]
            try:
                # Perform face swap
                swapped_frame = swapper.get(frame, face_in_frame, default_face, paste_back=True)
            except Exception as e:
                print(f"Error during face swap: {e}")
                swapped_frame = frame
        else:
            swapped_frame = frame

        _, buffer = cv2.imencode('.jpg', swapped_frame)
        processed_img_data = base64.b64encode(buffer).decode('utf-8')
        return {'image': 'data:image/jpeg;base64,' + processed_img_data}
    except Exception as e:
        print(f"Error processing frame: {e}")
        return {'error': str(e)}, 500


@app.route('/get_current_face')
@require_auth
def get_current_face():
    return send_file(default_img_path, mimetype='image/jpeg')


@app.route('/upload_face', methods=['POST'])
@require_auth
def upload_face():
    if 'face' not in request.files:
        return jsonify({'success': False, 'message': 'No file part'})

    file = request.files['face']
    if file.filename == '':
        return jsonify({'success': False, 'message': 'No selected file'})

    # Create a backup of the original default face if it doesn't exist
    backup_path = 'assets/face.jpg'
    if not os.path.exists(backup_path):
        shutil.copy2(default_img_path, backup_path)

    # Save the new face (always as face.jpg)
    file_path = os.path.join('assets', 'face.jpg')
    file.save(file_path)

    # Reload the face for face swapping
    global default_img, default_face
    default_img = cv2.imread(default_img_path)
    faces_in_default_img = app_model.get(default_img)

    if len(faces_in_default_img) == 0:
        # If no face detected, restore the backup
        shutil.copy2(backup_path, default_img_path)
        default_img = cv2.imread(default_img_path)
        faces_in_default_img = app_model.get(default_img)
        return jsonify({'success': False, 'message': 'No face detected in the uploaded image'})

    default_face = faces_in_default_img[0]
    return jsonify({'success': True})


@app.route('/reset_face', methods=['POST'])
@require_auth
def reset_face():
    backup_path = 'assets/face.jpg'
    if os.path.exists(backup_path):
        # Restore the original face
        shutil.copy2(backup_path, default_img_path)

        # Reload the face for face swapping
        global default_img, default_face
        default_img = cv2.imread(default_img_path)
        faces_in_default_img = app_model.get(default_img)
        default_face = faces_in_default_img[0]
        return jsonify({'success': True})
    else:
        return jsonify({'success': False, 'message': 'No backup found'})


@app.errorhandler(401)
def unauthorized(error):
    return jsonify({'error': 'Unauthorized access. Valid key required.'}), 401


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)