api/recognize.py

"""
api/recognize.py — POST /api/recognitions

Luồng nhận diện đầy đủ theo workflow diagram:
    1. Nhận ảnh → lưu raw_image vào MongoDB
    2. AI detect bounding boxes
    3. Với mỗi face:
       a. Crop face → lưu face_crop vào MongoDB (metadata sơ bộ)
       b. InsightFace tạo embedding 512D
       c. Gọi Supabase RPC match_face (fallback → RAM cosine similarity)
       d. MATCHED  → update visits/persons/recognition_logs, cập nhật face_crop links
       e. UNKNOWN  → ghi recognition_logs (person_id=NULL), lưu unknown_faces + recognition_payloads
    4. Trả về annotated_image + results[]
"""
import logging
import time

from flask import Blueprint, jsonify, request

import config
from face_model.analyzer import face_analyzer
from storage.mongo_storage import mongo_storage
from storage.pg_storage import pg_storage
from utils.image import (
    annotate_frame,
    crop_face,
    decode_image,
    get_image_dimensions,
    image_to_base64,
    image_to_bytes,
)

logger = logging.getLogger("api.recognize")

bp = Blueprint("recognize", __name__, url_prefix="/api")

# Màu bbox trên ảnh annotated
COLOR_MATCHED = (0, 220, 80)   # Xanh lá — đã nhận diện
COLOR_UNKNOWN = (0, 80, 255)   # Xanh dương — chưa biết
COLOR_NO_DB   = (0, 165, 255)  # Cam — DB rỗng


def _do_recognition(file_bytes: bytes) -> tuple[dict, int]:
    """Core logic nhận diện — tái sử dụng cho cả endpoint mới và legacy."""
    start_ms = time.time() * 1000

    # ── 1. Decode ảnh ────────────────────────────────────────────────────
    img = decode_image(file_bytes)
    if img is None:
        return {"error": {
            "code": "INVALID_IMAGE",
            "message": "Không đọc được ảnh.",
            "status": 400,
        }}, 400

    width, height = get_image_dimensions(img)

    # ── 2. Lưu raw image vào MongoDB ─────────────────────────────────────
    raw_image_id = mongo_storage.save_raw_image(
        img_bytes=file_bytes,
        source_type="webcam",
        width=width,
        height=height,
        fmt="jpeg",
    )

    # ── 3. Detect khuôn mặt ──────────────────────────────────────────────
    faces = face_analyzer.get_faces(img)
    if not faces:
        return {
            "total_detected": 0,
            "results": [],
            "annotated_image": image_to_base64(img),
            "raw_image_id": raw_image_id,
        }, 200

    results = []
    frame = img.copy()

    for idx, face in enumerate(faces):
        bbox_arr  = face.bbox.astype(int)
        embedding = face.normed_embedding
        det_score = float(face.det_score) if hasattr(face, "det_score") else 0.0

        # ── 3a. Crop khuôn mặt ───────────────────────────────────────────
        face_img, bbox_coords = crop_face(frame, bbox_arr, padding=20)
        crop_bytes = image_to_bytes(face_img)

        # ── 3b. Lưu face_crop sơ bộ (chưa có person_id/visit_id) ─────────
        crop_id = mongo_storage.save_face_crop(
            img_bytes=crop_bytes,
            raw_image_id=raw_image_id,
            bbox=bbox_coords,
            image_type="recognition",
            confidence=det_score,
            quality_score=det_score,
            model_name=config.MODEL_NAME,
        )

        # ── 3c. Tìm khuôn mặt khớp ───────────────────────────────────────
        emb_list = embedding.tolist()
        threshold = config.SIMILARITY_THRESHOLD

        # Thử RPC trước, nếu thất bại → fallback RAM
        rpc_result = pg_storage.match_face_rpc(emb_list, threshold)

        if rpc_result:
            max_sim  = float(rpc_result.get("similarity", 0.0))
            person_id   = rpc_result.get("person_id", "")
            matched_name = rpc_result.get("full_name", "")
            embedding_id = rpc_result.get("embedding_id", "")
            best_idx = -1
            logger.info(f"[RPC Match] {matched_name} sim={max_sim:.3f}")
        else:
            max_sim, best_idx = face_analyzer.find_match_ram(embedding, threshold)
            if best_idx >= 0:
                person_id    = face_analyzer.known_person_ids[best_idx]
                matched_name = face_analyzer.known_names[best_idx]
                embedding_id = face_analyzer.known_embedding_ids[best_idx]
            else:
                person_id    = ""
                matched_name = ""
                embedding_id = ""

        matched = bool(person_id) and max_sim >= threshold

        if matched:
            # ── 3d. MATCHED ──────────────────────────────────────────────
            ui_color = "#{:02X}{:02X}{:02X}".format(*COLOR_MATCHED[::-1])

            try:
                visit_id, log_id = pg_storage.record_matched_visit(
                    person_id=person_id,
                    embedding_id=embedding_id,
                    face_crop_mongo_id=crop_id,
                    raw_image_mongo_id=raw_image_id,
                    similarity=max_sim,
                    bbox=bbox_coords,
                    ui_color=ui_color,
                    ui_track_index=idx,
                )
                # Cập nhật links trong face_crop
                if crop_id:
                    mongo_storage.update_face_crop_links(
                        crop_id=crop_id,
                        person_id=person_id,
                        visit_id=visit_id,
                        recognition_log_id=log_id,
                    )
            except Exception as e:
                logger.error(f"[recognize] record_matched_visit failed: {e}")
                visit_id, log_id = "", ""

            annotate_frame(
                frame, bbox_arr,
                f"{matched_name} {max_sim * 100:.1f}%",
                COLOR_MATCHED,
            )
            results.append({
                "face_index":        idx + 1,
                "bbox":              bbox_arr.tolist(),
                "status":            "matched",
                "name":              matched_name,
                "person_id":         person_id,
                "similarity":        round(max_sim * 100, 2),
                "crop_id":           crop_id,
                "matched_image_url": f"/api/face-crops/{crop_id}" if crop_id else None,
            })

        else:
            # ── 3e. UNKNOWN ──────────────────────────────────────────────
            reason = "no_db" if face_analyzer.total == 0 else "no_match"
            ui_color_hex = "#{:02X}{:02X}{:02X}".format(
                *(COLOR_NO_DB[::-1] if face_analyzer.total == 0 else COLOR_UNKNOWN[::-1])
            )

            try:
                log_id = pg_storage.record_unknown_log(
                    face_crop_mongo_id=crop_id,
                    raw_image_mongo_id=raw_image_id,
                    similarity=max_sim,
                    bbox=bbox_coords,
                    ui_color=ui_color_hex,
                    ui_track_index=idx,
                )
            except Exception as e:
                logger.error(f"[recognize] record_unknown_log failed: {e}")
                log_id = ""

            # Lưu unknown_face vào MongoDB
            unknown_id = mongo_storage.save_unknown_face(
                face_crop_id=crop_id,
                raw_image_id=raw_image_id,
                recognition_log_id=log_id,
                reason=reason,
                confidence=det_score,
                similarity_distance=float(1.0 - max_sim),
            )

            # Lưu recognition_payload để debug
            elapsed_ms = time.time() * 1000 - start_ms
            mongo_storage.save_recognition_payload(
                recognition_log_id=log_id,
                raw_image_id=raw_image_id,
                face_crop_id=crop_id,
                model_name=config.MODEL_NAME,
                model_version="1.0.0",
                request_payload={
                    "face_index": idx,
                    "bbox": bbox_coords,
                    "det_score": det_score,
                    "threshold": threshold,
                    "match_method": "rpc" if rpc_result is not None else "ram",
                },
                response_payload={
                    "max_similarity": max_sim,
                    "reason": reason,
                    "unknown_id": unknown_id,
                },
                runtime_ms=elapsed_ms,
            )

            color = COLOR_NO_DB if face_analyzer.total == 0 else COLOR_UNKNOWN
            label = "Empty DB" if face_analyzer.total == 0 else "Unknown"
            annotate_frame(frame, bbox_arr, label, color)

            results.append({
                "face_index":      idx + 1,
                "bbox":            bbox_arr.tolist(),
                "status":          "unknown",
                "name":            None,
                "person_id":       None,
                "similarity":      round(max_sim * 100, 2),
                "reason":          reason,
                "crop_id":         crop_id,
                "unknown_face_id": unknown_id,
            })

    return {
        "total_detected":  len(faces),
        "results":         results,
        "annotated_image": image_to_base64(frame),
        "raw_image_id":    raw_image_id,
    }, 200


@bp.post("/recognitions")
def recognitions():
    """
    POST /api/recognitions
    Body: multipart/form-data với field 'file' chứa ảnh JPEG/PNG
    """
    if "file" not in request.files:
        return jsonify({"error": {
            "code": "FILE_REQUIRED",
            "message": "Không tìm thấy file ảnh trong request.",
            "status": 400,
        }}), 400

    file = request.files["file"]
    if file.filename == "":
        return jsonify({"error": {
            "code": "EMPTY_FILENAME",
            "message": "Tên file rỗng.",
            "status": 400,
        }}), 400

    try:
        file_bytes = file.read()
        data, status_code = _do_recognition(file_bytes)
        return jsonify(data), status_code
    except Exception as e:
        logger.error(f"[recognitions] Unexpected error: {e}", exc_info=True)
        return jsonify({"error": {
            "code": "RECOGNITION_FAILED",
            "message": f"Lỗi trong quá trình nhận diện: {e}",
            "status": 500,
        }}), 500


@bp.post("/recognize")
def recognize_legacy():
    """
    Legacy alias → POST /api/recognitions (backward compatibility với frontend).
    Response format giữ nguyên để frontend không cần sửa.
    """
    if "file" not in request.files:
        return jsonify({"detail": "Không tìm thấy file ảnh để nhận diện."}), 400

    file = request.files["file"]
    try:
        file_bytes = file.read()
        data, status_code = _do_recognition(file_bytes)

        if status_code != 200:
            error = data.get("error", {})
            return jsonify({"detail": error.get("message", "Lỗi.")}), status_code

        # Transform response về format cũ mà frontend đang dùng
        legacy_results = []
        for r in data.get("results", []):
            legacy_r = {
                "face_index":  r["face_index"],
                "bbox":        r["bbox"],
                "status":      r["status"],
                "name":        r.get("name") or f"Unknown_{r['face_index']}",
                "person_id":   r.get("person_id") or "",
                "similarity":  r.get("similarity", 0),
            }
            if r["status"] == "matched":
                legacy_r["matched_image_url"] = r.get("matched_image_url", "")
            else:
                legacy_r["saved_as"]        = f"unknown_{r['face_index']}.jpg"
                legacy_r["saved_image_url"] = f"/api/face-crops/{r.get('crop_id', '')}"
            legacy_results.append(legacy_r)

        return jsonify({
            "total_detected":  data["total_detected"],
            "results":         legacy_results,
            "annotated_image": data["annotated_image"],
        }), 200

    except Exception as e:
        logger.error(f"[recognize_legacy] Unexpected error: {e}", exc_info=True)
        return jsonify({"detail": f"Lỗi trong quá trình nhận diện: {e}"}), 500