"""
face_engine.py
--------------
The face-recognition core, rebuilt around embeddings + similarity matching
instead of a retrained softmax classifier. See the README for the full
rationale; short version:

    OLD: photos -> train a fresh N-class classifier -> predict a class index
         (needs 2+ students, full retrain to add anyone, poor stranger-rejection)

    NEW: photos -> CNN embedding (a 1280-d "fingerprint" vector) -> stored in
         a small gallery file -> new faces are matched by cosine similarity
         against every stored vector
         (works with 1 student, registering someone is instant, strangers are
         rejected by threshold + margin instead of forced into a class)

Pipeline for a single photo:
  1. DETECT   - find the face region. Haar Cascade runs first (fast); if it
                finds nothing, MTCNN (a small CNN-based detector) is tried as
                a fallback since it handles angled/harder faces better.
  2. QUALITY  - reject/warn on faces that are too small, blurry (Laplacian
                variance), or too dark/bright, using OpenCV metrics -- this
                catches bad registration photos before they ever hurt
                recognition accuracy.
  3. PREPROCESS - resize to 96x96, normalize for MobileNetV2.
  4. EMBED    - MobileNetV2 (frozen, ImageNet weights, pooling='avg') maps the
                face to a 1280-d vector. No training involved -- this is pure
                feature extraction, which is why registration is instant.
  5. MATCH    - the new embedding is compared via cosine similarity against
                every embedding in the gallery (static/... no, DATA_DIR/embeddings/
                embeddings.json). The closest student wins IF (a) the
                similarity clears MATCH_THRESHOLD and (b) it beats the
                second-best candidate by at least MATCH_MARGIN -- the margin
                check is what catches "two plausible but wrong" matches that
                a bare threshold would let through.
"""

import os
import json
import tempfile
import shutil
import numpy as np
import config

try:
    import cv2
    _haar_cascade = cv2.CascadeClassifier(
        cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
    )
except ImportError:
    cv2 = None
    _haar_cascade = None
    print("[face_engine] OpenCV not available - install with 'pip install opencv-python-headless'")

_mtcnn_detector = None


def _get_mtcnn():
    global _mtcnn_detector
    if _mtcnn_detector is None:
        try:
            from mtcnn import MTCNN
            _mtcnn_detector = MTCNN()
        except ImportError:
            print("[face_engine] MTCNN not available - install with 'pip install mtcnn'")
            return None
    return _mtcnn_detector


def detect_face(image_bgr):
    """
    Finds the largest/most confident face in a BGR image and returns the
    cropped face region (BGR). Returns None if no face could be found by
    either detector.
    """
    if cv2 is None:
        print("[face_engine] OpenCV not available, cannot detect faces")
        return None

    if _haar_cascade is None or _haar_cascade.empty():
        # Haar unavailable -- try MTCNN exclusively
        mtcnn = _get_mtcnn()
        if mtcnn is not None:
            try:
                image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
                results = mtcnn.detect_faces(image_rgb)
                if results:
                    best = max(results, key=lambda r: r["box"][2] * r["box"][3])
                    x, y, w, h = best["box"]
                    x, y = max(0, x), max(0, y)
                    face = image_bgr[y:y + h, x:x + w]
                    if face.size > 0:
                        return face
            except Exception as e:
                print(f"[face_engine] MTCNN detection failed: {e}")
        return None

    try:
        gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
        faces = _haar_cascade.detectMultiScale(
            gray, scaleFactor=1.05, minNeighbors=3, minSize=(30, 30)
        )

        if len(faces) > 0:
            x, y, w, h = max(faces, key=lambda box: box[2] * box[3])
            return image_bgr[y:y + h, x:x + w]

        # Haar found nothing -- try MTCNN fallback
        mtcnn = _get_mtcnn()
        if mtcnn is not None:
            try:
                image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
                results = mtcnn.detect_faces(image_rgb)
                if results:
                    best = max(results, key=lambda r: r["box"][2] * r["box"][3])
                    x, y, w, h = best["box"]
                    x, y = max(0, x), max(0, y)
                    face = image_bgr[y:y + h, x:x + w]
                    if face.size > 0:
                        return face
            except Exception as e:
                print(f"[face_engine] MTCNN fallback failed: {e}")

    except Exception as e:
        print(f"[face_engine] Error in detect_face: {e}")
        return None

    return None


def assess_quality(face_bgr):
    """
    Runs cheap, fast heuristics on a cropped face and returns a list of
    human-readable warning strings (empty list = looks good).
    """
    warnings = []
    gray = cv2.cvtColor(face_bgr, cv2.COLOR_BGR2GRAY)
    h, w = gray.shape[:2]

    if min(h, w) < config.MIN_FACE_SIZE:
        warnings.append(f"Face looks small in frame ({w}x{h}px) — try moving closer.")

    blur_score = cv2.Laplacian(gray, cv2.CV_64F).var()
    if blur_score < config.BLUR_THRESHOLD:
        warnings.append("Image looks blurry — hold still and make sure the camera is focused.")

    brightness = float(np.mean(gray))
    if brightness < config.MIN_BRIGHTNESS:
        warnings.append("Image is quite dark — try better lighting.")
    elif brightness > config.MAX_BRIGHTNESS:
        warnings.append("Image is overexposed — reduce direct light or glare.")

    return warnings


# ---------------------------------------------------------------------------
# Preprocessing + embedding extraction
# ---------------------------------------------------------------------------
def preprocess_face(face_bgr):
    """Resize to IMG_SIZE and normalize the way MobileNetV2 expects."""
    from keras.applications.mobilenet_v2 import preprocess_input

    if face_bgr is None:
        raise ValueError("preprocess_face(): received None")

    if not isinstance(face_bgr, np.ndarray):
        raise TypeError(f"preprocess_face(): expected numpy array, got {type(face_bgr)}")

    if face_bgr.ndim != 3 or face_bgr.shape[2] != 3:
        raise ValueError(f"preprocess_face(): expected HxWx3 BGR image, got shape {face_bgr.shape}")

    h, w = face_bgr.shape[:2]
    if h < 2 or w < 2:
        raise ValueError(f"preprocess_face(): crop too small with shape {face_bgr.shape}")

    face_resized = cv2.resize(
        face_bgr, (config.IMG_SIZE, config.IMG_SIZE), interpolation=cv2.INTER_LINEAR
    )
    face_rgb = cv2.cvtColor(face_resized, cv2.COLOR_BGR2RGB)

    # Add batch dimension: (96, 96, 3) -> (1, 96, 96, 3)
    face_rgb = np.expand_dims(face_rgb, axis=0)

    out = preprocess_input(face_rgb)

    if not (isinstance(out, np.ndarray) and out.shape == (1, config.IMG_SIZE, config.IMG_SIZE, 3)):
        raise ValueError(f"preprocess_face(): unexpected output shape {getattr(out, 'shape', None)}")

    return out


_embedder = None


def _get_embedder():
    """
    Lazy-loads MobileNetV2 as a pure feature extractor.
    """
    global _embedder
    if _embedder is None:
        from keras.applications.mobilenet_v2 import MobileNetV2
        _embedder = MobileNetV2(
            input_shape=(config.IMG_SIZE, config.IMG_SIZE, 3),
            include_top=False,
            weights="imagenet",
            pooling="avg"
        )
    return _embedder


def warm_up_models():
    """Warm up all models at startup to reduce latency during first use."""
    print("[face_engine] Warming up models...")
    try:
        embedder = _get_embedder()
        # Lightweight warm-up: single dummy prediction
        dummy = np.zeros((1, config.IMG_SIZE, config.IMG_SIZE, 3), dtype="float32")
        embedder.predict(dummy, verbose=0)
        print("[face_engine] Embedder warmed up.")
    except Exception as e:
        print(f"[face_engine] Embedder warm-up failed: {e}")

    try:
        mtcnn = _get_mtcnn()
        if mtcnn:
            print("[face_engine] MTCNN warmed up.")
    except Exception as e:
        print(f"[face_engine] MTCNN warm-up failed: {e}")


def compute_embedding(face_bgr):
    """Returns an L2-normalized 1280-d embedding vector for a cropped face."""
    embedder = _get_embedder()
    face_array = preprocess_face(face_bgr)

    if face_array.ndim != 4 or face_array.shape[1:] != (config.IMG_SIZE, config.IMG_SIZE, 3):
        raise ValueError(f"compute_embedding(): unexpected batch shape {face_array.shape}")

    raw = embedder.predict(face_array, verbose=0)[0]
    norm = np.linalg.norm(raw)
    return (raw / norm) if norm > 0 else raw


def cosine_similarity(a, b):
    """Dot product of two already-L2-normalized vectors == cosine similarity."""
    return float(np.dot(a, b))


# ---------------------------------------------------------------------------
# Gallery persistence  (ATOMIC writes to prevent corruption on crash/OOM)
# ---------------------------------------------------------------------------
def load_gallery():
    """Returns {student_id: [embedding, embedding, ...]} as numpy arrays."""
    if not os.path.exists(config.EMBEDDINGS_PATH):
        return {}
    with open(config.EMBEDDINGS_PATH) as f:
        raw = json.load(f)
    return {sid: [np.array(e, dtype="float32") for e in embeddings] for sid, embeddings in raw.items()}


def save_gallery(gallery):
    """Atomically write gallery JSON to prevent corruption if process is killed mid-write."""
    config.ensure_directories()
    serializable = {sid: [e.tolist() for e in embeddings] for sid, embeddings in gallery.items()}

    # Write to temp file in same directory, then atomic rename
    dir_name = os.path.dirname(config.EMBEDDINGS_PATH)
    fd, tmp_path = tempfile.mkstemp(dir=dir_name, suffix=".tmp")
    try:
        with os.fdopen(fd, "w") as f:
            json.dump(serializable, f)
        shutil.move(tmp_path, config.EMBEDDINGS_PATH)
    except Exception:
        # Clean up temp file on failure
        try:
            os.remove(tmp_path)
        except OSError:
            pass
        raise


def add_photo_to_gallery(student_id, face_bgr):
    """Computes and stores one more reference embedding for a student."""
    try:
        embedding = compute_embedding(face_bgr)
    except Exception as e:
        print(f"[face_engine] compute_embedding failed for {student_id}: {e}")
        raise RuntimeError(f"Could not compute face embedding: {e}")

    gallery = load_gallery()
    gallery.setdefault(student_id, []).append(embedding)
    save_gallery(gallery)
    return embedding


def remove_student_from_gallery(student_id):
    gallery = load_gallery()
    if student_id in gallery:
        del gallery[student_id]
        save_gallery(gallery)


def reindex_gallery():
    """
    Rebuilds embeddings.json from scratch by re-reading every photo under
    DATASET_DIR/<student_id>/*. Useful after bulk-importing photos directly
    onto disk, or if the embedding model ever changes and old vectors need
    recomputing.
    """
    gallery = {}
    processed_images = 0
    skipped_images = 0

    if os.path.isdir(config.DATASET_DIR):
        for student_id in sorted(os.listdir(config.DATASET_DIR)):
            folder = os.path.join(config.DATASET_DIR, student_id)
            if not os.path.isdir(folder):
                continue

            embeddings = []
            try:
                for filename in sorted(os.listdir(folder)):
                    filepath = os.path.join(folder, filename)
                    try:
                        image_bgr = cv2.imread(filepath)
                        if image_bgr is None:
                            skipped_images += 1
                            continue
                        face = detect_face(image_bgr)
                        if face is None:
                            skipped_images += 1
                            continue
                        embeddings.append(compute_embedding(face))
                        processed_images += 1
                    except Exception as e:
                        print(f"[face_engine] Error processing image {filepath}: {e}")
                        skipped_images += 1
            except PermissionError:
                print(f"[face_engine] Permission denied accessing folder {folder}")
                continue

            if embeddings:
                gallery[student_id] = embeddings

    save_gallery(gallery)
    return {
        "students_indexed": len(gallery),
        "images_processed": processed_images,
        "images_skipped": skipped_images
    }


# ---------------------------------------------------------------------------
# Recognition
# ---------------------------------------------------------------------------
def match_face(image_bgr):
    """
    Full pipeline for an attendance check: detect -> embed -> compare.

    Returns a dict:
        {"student_id": "STU001", "confidence": 0.81}                    on match
        {"student_id": None, "confidence": 0.0, "reason": "..."}        otherwise
    """
    face = detect_face(image_bgr)
    if face is None:
        return {"student_id": None, "confidence": 0.0, "reason": "No face detected in the image."}

    gallery = load_gallery()
    if not gallery:
        return {"student_id": None, "confidence": 0.0, "reason": "No students registered yet."}

    try:
        query_embedding = compute_embedding(face)
    except Exception as e:
        print(f"[face_engine] Failed to compute embedding: {e}")
        return {"student_id": None, "confidence": 0.0, "reason": f"Could not compute face embedding: {e}"}

    # Flatten all embeddings with their student IDs
    student_ids = []
    all_embeddings = []
    for sid, embeddings in gallery.items():
        for emb in embeddings:
            student_ids.append(sid)
            all_embeddings.append(emb)

    if not all_embeddings:
        return {"student_id": None, "confidence": 0.0, "reason": "No valid face embeddings to compare against."}

    all_embeddings = np.array(all_embeddings)
    query_embedding = np.array(query_embedding)

    # Compute similarities in batch
    similarities = np.dot(all_embeddings, query_embedding)

    # Group by student and find best score per student
    best_per_student = {}
    for i, sid in enumerate(student_ids):
        sim = similarities[i]
        if sid not in best_per_student or best_per_student[sid] < sim:
            best_per_student[sid] = sim

    ranked = sorted(best_per_student.items(), key=lambda item: item[1], reverse=True)
    best_id, best_score = ranked[0]
    second_score = ranked[1][1] if len(ranked) > 1 else -1.0

    if best_score < config.MATCH_THRESHOLD:
        return {
            "student_id": None,
            "confidence": best_score,
            "reason": f"Face not recognized — confidence {best_score:.2f} is below threshold {config.MATCH_THRESHOLD:.2f}."
        }

    if (best_score - second_score) < config.MATCH_MARGIN and len(ranked) > 1:
        return {
            "student_id": None,
            "confidence": best_score,
            "reason": "Match too close between two students — please retake the photo."
        }

    return {"student_id": best_id, "confidence": best_score}