import cv2
import numpy as np
import onnxruntime as ort
import os
import urllib.request
import json

class Authenticator:
    """
    Module: Biometric Security Engine (ONNX Powered)
    Handles military-grade facial verification for Aelion OS.
    """
    def __init__(self):
        self.current_dir = os.path.dirname(os.path.abspath(__file__))
        self.model_path = os.path.join(self.current_dir, "mobilefacenet.onnx")
        self.db_path = os.path.join(self.current_dir, "aelion_identities.json")

        # --- THE AUTO-DOWNLOADER ---
        if not os.path.exists(self.model_path):
            self._download_model()

        print("[Aelion Vision] Booting Biometric Security Engine...")
        try:
            self.session = ort.InferenceSession(self.model_path)
        except Exception as e:
            print(f"[!] Critical Error: Could not load biometric weights. {e}")
            self.session = None

    def _download_model(self):
        """Grabs the MobileFaceNet ONNX weights."""
        print("[Aelion Vision] Initializing Security Protocols for the first time...")
        print("[Aelion Vision] Downloading MobileFaceNet Weights (~5MB)...")
        
        # TODO: Swap this with your actual Aelion API / Server URL
        url = "https://your-aelion-server.com/models/mobilefacenet.onnx" 
        
        try:
            urllib.request.urlretrieve(url, self.model_path)
            print("[Aelion Vision] Download Complete! Biometrics online.")
        except Exception as e:
            print(f"[!] Critical Error: Failed to download the AI model. {e}")

    def _get_face_vector(self, frame, face_box):
        """Extracts the 128-dimensional mathematical fingerprint of a face."""
        if self.session is None or frame is None:
            return None

        try:
            x, y, w, h = face_box
            
            # Failsafe bounds
            h_f, w_f = frame.shape[:2]
            x, y = max(0, x), max(0, y)
            w, h = min(w_f - x, w), min(h_f - y, h)
            
            face_img = frame[y:y+h, x:x+w]
            if face_img.size == 0: return None

            # 1. MobileFaceNet expects exactly 112x112 RGB images
            face_resized = cv2.resize(face_img, (112, 112))
            face_rgb = cv2.cvtColor(face_resized, cv2.COLOR_BGR2RGB)

            # 2. Normalize pixel values to range [-1, 1]
            face_normalized = (face_rgb.astype(np.float32) - 127.5) / 128.0

            # 3. Reorder to NCHW format
            input_data = np.transpose(face_normalized, (2, 0, 1))
            input_data = np.expand_dims(input_data, axis=0)

            # 4. Run the AI Brain
            inputs = {self.session.get_inputs()[0].name: input_data}
            outputs = self.session.run(None, inputs)
            
            # Return the raw 128-D vector (flattened)
            return outputs[0].flatten()

        except Exception as e:
            print(f"[!] Vector extraction failed: {e}")
            return None

    def register_owner(self, frame, face_box, user_id="Admin"):
        """Saves your face vector to a local JSON file to lock the system."""
        vector = self._get_face_vector(frame, face_box)
        if vector is None:
            print("[Aelion Security] Could not detect face clearly. Try again.")
            return False

        # Load existing DB or create new one
        db = {}
        if os.path.exists(self.db_path):
            with open(self.db_path, "r") as f:
                db = json.load(f)

        # Save the vector (must convert numpy array to standard python list for JSON)
        db[user_id] = vector.tolist()
        
        with open(self.db_path, "w") as f:
            json.dump(db, f)
            
        print(f"[Aelion Security] Identity registered for: {user_id}")
        return True

    def verify_user(self, frame, face_box, user_id="Admin", tolerance=0.55):
        """
        Compares the face on camera to the saved JSON vector.
        Uses Cosine Similarity. Higher score = closer match.
        """
        if not os.path.exists(self.db_path):
            return "UNREGISTERED"

        with open(self.db_path, "r") as f:
            db = json.load(f)

        if user_id not in db:
            return "UNKNOWN_ID"

        current_vector = self._get_face_vector(frame, face_box)
        if current_vector is None: return "SCANNING..."

        saved_vector = np.array(db[user_id])

        # --- THE MATH: Cosine Similarity ---
        dot_product = np.dot(current_vector, saved_vector)
        norm_curr = np.linalg.norm(current_vector)
        norm_saved = np.linalg.norm(saved_vector)
        
        similarity = dot_product / (norm_curr * norm_saved)

        # If similarity is above the tolerance, it's you.
        if similarity > tolerance:
            return "ACCESS_GRANTED"
        else:
            return "ACCESS_DENIED"