import base64
import os
import tempfile

# Use temporary directory or current working directory for app data
# This avoids permission issues in containerized environments
# Check if running in container with pre-created directories
if os.path.exists('/tmp/app_data') and os.access('/tmp/app_data', os.W_OK):
    app_data_path = '/tmp/app_data'
    print("Using pre-created /tmp/app_data directory")
else:
    try:
        # Try to use /tmp first (usually writable in containers)
        app_data_path = '/tmp/app_data'
        os.makedirs(app_data_path, exist_ok=True)
    except PermissionError:
        # Fallback to current directory or temp directory
        app_data_path = os.path.join(os.getcwd(), 'app_data')
        try:
            os.makedirs(app_data_path, exist_ok=True)
        except PermissionError:
            # Last resort: use system temp directory
            app_data_path = tempfile.mkdtemp(prefix='app_data_')

# Make sure the necessary directories exist and have proper permissions

import cv2
import numpy as np
from flask import Flask, request, jsonify
from mtcnn.mtcnn import MTCNN
from keras_facenet import FaceNet
from sklearn.metrics.pairwise import cosine_similarity
from flask_cors import CORS
from pymongo import MongoClient
from pymongo.server_api import ServerApi

# Initialize MongoDB connection
client = MongoClient('mongodb+srv://nanduvinay719:76qqKRX4zC97yQun@travis.744fuyn.mongodb.net/?retryWrites=true&w=majority&appName=travis', server_api=ServerApi('1'))
if client:
    print("Connected to MongoDB")
db = client["travis"]
mongo = db["travis_face_data"]
if "travis_face_data" not in db.list_collection_names():
    db.create_collection("travis_face_data")

app = Flask(__name__)
CORS(app)

# Initialize MTCNN detector and FaceNet model
detector = MTCNN()
embedder = FaceNet()


def cosine(embedding1, embedding2):
    dot_product = np.dot(embedding1, embedding2)
    norm1 = np.linalg.norm(embedding1)
    norm2 = np.linalg.norm(embedding2)
    similarity = dot_product / (norm1 * norm2)
    return similarity

def reload_embeddings():
    global face_data, labels, names
    face_data, labels, names = load_embeddings_from_db()

@app.route("/")
def home():
    return {"message": "Travis Login API is running!"}


@app.route('/login', methods=['POST'])
def recognizeLogin():
    try:
        if 'image' not in request.files:
            return jsonify({"error": "No image provided"}), 400

        file = request.files['image']
        if not file:
            return jsonify({"error": "Invalid file"}), 400

        image_data = file.read()
        image_array = np.frombuffer(image_data, np.uint8)
        image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)   
        if image is None:
            return jsonify({"error": "Invalid image"}), 400

        results = recognize_faces_in_image(image)
        print(results[0])
        if results[0]['name'] != 'unknown':
            return jsonify({"name": results[0]['name'], "probability": results[0]['probability']}), 200
        else:
            return jsonify({'name': "user not recognised"}), 401
    except Exception as e:
        print(f"Error in login: {str(e)}")
        return jsonify({"error": "Internal server error"}), 500

@app.route('/register', methods=['POST'])
def register():
    username = request.form['username']
    phoneno = request.form["phoneNumber"]
    email = request.form['email']
    facenet_embeddings = []
    stored_image = None  # To store the first grayscale image

    print(username)
    # Check if user already exists
    existing_user = mongo.find_one({"username": username})
    if existing_user:
        return jsonify({"error": f"User '{username}' already exists"}), 400

    # Process uploaded images
    for i in range(5):  # Expecting 5 images
        try:
            image_file = request.files[f'image{i}']
        except KeyError:
            return jsonify({"error": f"Missing image{i} in the request"}), 400

        image_data = image_file.read()
        image_array = np.frombuffer(image_data, np.uint8)
        image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)

        # Face detection using MTCNN for FaceNet
        mtcnn_faces = detector.detect_faces(image)
        if mtcnn_faces:
            # Get the first detected face for FaceNet embedding
            x, y, w, h = mtcnn_faces[0]['box']
            x, y = max(0, x), max(0, y)
            cropped_face = cv2.resize(image[y:y+h, x:x+w], (160, 160))
            rgb_face = cv2.cvtColor(cropped_face, cv2.COLOR_BGR2RGB)

            # Get FaceNet embedding
            facenet_embedding = embedder.embeddings(np.expand_dims(rgb_face, axis=0)).flatten()
            facenet_embeddings.append(facenet_embedding)
        # Face detection using Haar Cascade for CNN
            if stored_image is None:
                _, buffer = cv2.imencode('.jpg', cv2.cvtColor(cropped_face, cv2.COLOR_BGR2GRAY))
                stored_image = base64.b64encode(buffer).decode('utf-8')

    if not facenet_embeddings:
        return jsonify({"error": "No valid faces detected in the uploaded images"}), 400

    # Calculate mean embeddings
    mean_facenet_embedding = np.mean(facenet_embeddings, axis=0).astype(float).tolist()
    
    # Save model weights with error handling
    
    # Create user data
    id = mongo.count_documents({}) + 1
    user_data = {
        'username': username,
        'phoneNumber': phoneno,
        'email': email, 
        'embeddings': mean_facenet_embedding,
        'stored_image': stored_image,
        'role': 'agent',
        'id': id
    }

    # Insert into MongoDB
    mongo.insert_one(user_data)
    
    reload_embeddings()

    return jsonify({"message": "User registered successfully!"}), 201

# Load embeddings from MongoDB for recognition

def load_embeddings_from_db():
    try:
        users = list(mongo.find({"role": "agent"}))  # Only get agent users
        face_data = []  # facenet embeddings
        labels = []  # id 1,2,3,..
        names = {}  # dict of id and username
        for user in users:
            if 'embeddings' in user:  # Only process users with embeddings
                face_data.append(user["embeddings"])
                labels.append(user['id'])
                names[user['id']] = user['username']

        print(f"Loaded {len(face_data)} user embeddings from database")
        return (face_data, labels, names) if face_data else ([], [], {})
    except Exception as e:
        print(f"Error loading embeddings from database: {e}")
        return [], [], {}

# Load face embeddings from MongoDB initially
face_data, labels, names = load_embeddings_from_db()

def recognize_faces_in_image(image):
    if len(face_data) == 0:
        return [{"name": "No registered faces", "probability": 0.0}]

    faces = detector.detect_faces(image)
    results = []
    for face in faces:
        x, y, width, height = face['box']
        cropped_face = cv2.resize(image[y:y+height, x:x+width], (160, 160))
        
        # Convert cropped face to RGB
        rgb_face = cv2.cvtColor(cropped_face, cv2.COLOR_BGR2RGB)
        embedding = embedder.embeddings(np.expand_dims(rgb_face, axis=0)).flatten()  # Use RGB face here

        # Compare with stored embeddings in MongoDB
        similarities = cosine_similarity([embedding], face_data)
        idx = np.argmax(similarities)
        best_match = similarities[0][idx]

        if best_match > 0.7:
            recognized_id = labels[idx] # Get the ObjectId
            recognized_name = names[recognized_id]  # Use ObjectId to get the username
            results.append({"name": recognized_name, "probability": float(best_match)})
        else:
            results.append({"name": "unknown", "probability": float(best_match)})
    return results

@app.route('/admin-login', methods=['POST'])
def admin_login():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')
    
    if not username or not password:
        return jsonify({"error": "Username and password are required"}), 400
        
    # Find admin user in database
    admin = mongo.find_one({
        "username": username,
        "role": "admin"
    })
    
    if not admin:
        return jsonify({"error": "Invalid credentials"}), 401
        
    # In a real application, you should hash passwords and compare hashes
    # For now, we'll just check if the password matches
    if admin.get('password') != password:
        return jsonify({"error": "Invalid credentials"}), 401
        
    return jsonify({
        "success": True,
        "username": username,
        "role": "admin"
    }), 200

if __name__ == '__main__':
    app.run(host='0.0.0.0',port=7860,debug=True)