import os
import cv2
import numpy as np
import onnxruntime as ort
import logging

logger = logging.getLogger(__name__)

class FaceRecognizer:
    def __init__(self, model_path="models/arcface.onnx"):
        self.model_path = model_path
        self.loaded = False
        self.session = None
        
        if os.path.exists(model_path):
            try:
                # Use CPUExecutionProvider for HF Spaces basic instance
                self.session = ort.InferenceSession(
                    model_path, 
                    providers=['CPUExecutionProvider']
                )
                self.loaded = True
                logger.info(f"ArcFace Recognizer loaded successfully from {model_path}")
            except Exception as e:
                logger.error(f"Error initializing ArcFace ONNX session: {e}")
        else:
            logger.warning(f"ArcFace recognition model file missing at {model_path}")

    def apply_clahe(self, face_crop):
        """
        Drawback Fix - Low Light Handling:
        Applies CLAHE histogram equalization on the luminance channel (LAB color space)
        to normalize brightness and contrast before facial recognition.
        """
        if face_crop is None or face_crop.size == 0:
            return face_crop
            
        try:
            # Convert BGR to LAB color space
            lab = cv2.cvtColor(face_crop, cv2.COLOR_BGR2LAB)
            l_channel, a_channel, b_channel = cv2.split(lab)
            
            # Apply CLAHE on L (Lightness) channel
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
            cl = clahe.apply(l_channel)
            
            # Merge and convert back to BGR
            limg = cv2.merge((cl, a_channel, b_channel))
            bgr_equalized = cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)
            return bgr_equalized
        except Exception as e:
            logger.error(f"Failed to apply CLAHE: {e}")
            return face_crop

    def preprocess(self, face_crop):
        """
        Resizes to 112x112, converts BGR to RGB, normalizes,
        transposes to CHW, and adds batch dimension.
        """
        # Resize to ArcFace input dimensions (112x112)
        img = cv2.resize(face_crop, (112, 112))
        
        # Convert BGR to RGB
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        # Normalize: (pixel - 127.5) / 128.0
        img = img.astype(np.float32)
        img = (img - 127.5) / 128.0
        
        # Transpose to CHW: (3, 112, 112)
        img = np.transpose(img, (2, 0, 1))
        
        # Add batch dimension: (1, 3, 112, 112)
        input_tensor = np.expand_dims(img, axis=0)
        return input_tensor

    def get_embedding(self, face_crop):
        """
        Generates L2 normalized 512-dimensional embedding for a face crop.
        Includes low-light CLAHE enhancement.
        """
        if not self.loaded or self.session is None:
            logger.warning("ArcFace model not loaded. Skipping embedding generation.")
            return None
            
        if face_crop is None or face_crop.size == 0:
            return None
            
        try:
            # Equalize lighting conditions
            enhanced_crop = self.apply_clahe(face_crop)
            
            # Preprocess to get float32 tensor
            input_tensor = self.preprocess(enhanced_crop)
            
            # Run inference
            outputs = self.session.run(
                None,
                {self.session.get_inputs()[0].name: input_tensor}
            )
            raw_embedding = outputs[0][0]
            
            # L2 normalize
            norm = np.linalg.norm(raw_embedding)
            normalized_embedding = raw_embedding / norm if norm > 0 else raw_embedding
            return normalized_embedding
        except Exception as e:
            logger.error(f"Error generating ArcFace embedding: {e}")
            return None

    def match_student(self, embedding, known_profiles, threshold=0.68):
        """
        Matches a face embedding against a dictionary of known student profiles.
        Enforces double-match ambiguity rejection.
        """
        if embedding is None or not known_profiles:
            return None, 0.0
            
        scores = []
        for roll_no, profile in known_profiles.items():
            # Embeddings are L2 normalized, so cosine similarity is just the dot product
            similarity = float(np.dot(embedding, profile["embedding"]))
            scores.append((roll_no, similarity))
            
        if not scores:
            return None, 0.0
            
        # Sort by similarity descending
        scores.sort(key=lambda x: x[1], reverse=True)
        top_match, top_score = scores[0]
        
        # Check if best match is below confidence threshold
        if top_score < threshold:
            return None, top_score
            
        # Drawback Fix - Wrong Match Prevention (Ambiguity check)
        if len(scores) > 1:
            second_match, second_score = scores[1]
            if (top_score - second_score) < 0.05:
                logger.warning(
                    f"Ambiguous match: top choice '{top_match}' (score {top_score:.3f}) "
                    f"is too close to second choice '{second_match}' (score {second_score:.3f}). "
                    f"Margin is {top_score - second_score:.3f} (< 0.05). Rejecting."
                )
                return "AMBIGUOUS", top_score
                
        return top_match, top_score