main.py

"""
Face Verification API - Python backend for face enrollment and authentication.
Replaces FaceIO with a self-hosted solution using face_recognition library.
"""

import base64
import io
import os
import json
import math
from typing import Optional, List
from contextlib import asynccontextmanager

import numpy as np
import face_recognition
from PIL import Image
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Configuration
FACE_MATCH_THRESHOLD = float(os.getenv("FACE_MATCH_THRESHOLD", "0.6"))
ALLOWED_ORIGINS = os.getenv("ALLOWED_ORIGINS", "*").split(",")


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan handler."""
    print("🚀 Face Verification API starting...")
    print(f"   Match threshold: {FACE_MATCH_THRESHOLD}")
    yield
    print("👋 Face Verification API shutting down...")


app = FastAPI(
    title="Face Verification API",
    description="Self-hosted face enrollment and authentication service",
    version="1.0.0",
    lifespan=lifespan
)

# Configure CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS if ALLOWED_ORIGINS[0] != "*" else ["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# Request/Response Models
class EnrollRequest(BaseModel):
    """Request model for face enrollment."""
    image_base64: str
    user_id: str
    metadata: Optional[dict] = None


class EnrollResponse(BaseModel):
    """Response model for face enrollment."""
    success: bool
    face_id: str
    face_embedding: List[float]
    message: str


class VerifyRequest(BaseModel):
    """Request model for face verification."""
    image_base64: str
    stored_embedding: List[float]
    user_id: Optional[str] = None


class VerifyResponse(BaseModel):
    """Response model for face verification."""
    success: bool
    match: bool
    confidence: float
    distance: float
    message: str


class HealthResponse(BaseModel):
    """Response model for health check."""
    status: str
    version: str


def decode_base64_image(base64_string: str) -> np.ndarray:
    """
    Decode a base64 image string to a numpy array for face_recognition.
    Supports both raw base64 and data URI format.
    """
    # Remove data URI prefix(es) if present
    # Handle cases like "data:image/jpeg;base64,data:image/png;base64,..."
    if "," in base64_string:
        base64_string = base64_string.split(",")[-1]
    
    # Clean up the string and fix padding
    base64_string = base64_string.strip()
    padding = len(base64_string) % 4
    if padding > 0:
        base64_string += "=" * (4 - padding)
    
    try:
        image_bytes = base64.b64decode(base64_string)
        image = Image.open(io.BytesIO(image_bytes))
        
        # Convert to RGB if necessary (face_recognition requires RGB)
        if image.mode != "RGB":
            image = image.convert("RGB")
        
        return np.array(image)
    except Exception as e:
        raise HTTPException(
            status_code=400,
            detail=f"Invalid image data: {str(e)}"
        )


def extract_face_encoding(image: np.ndarray) -> np.ndarray:
    """
    Extract face encoding from an image.
    Returns the 128-dimensional face encoding.
    """
    # Detect face locations
    face_locations = face_recognition.face_locations(image, model="hog")
    
    if len(face_locations) == 0:
        raise HTTPException(
            status_code=400,
            detail="No face detected in the image. Please ensure your face is clearly visible."
        )
    
    if len(face_locations) > 1:
        raise HTTPException(
            status_code=400,
            detail="Multiple faces detected. Please ensure only one face is visible."
        )
    
    # Extract face encoding
    face_encodings = face_recognition.face_encodings(image, face_locations)
    
    if len(face_encodings) == 0:
        raise HTTPException(
            status_code=400,
            detail="Could not extract face features. Please try again with better lighting."
        )
    
    return face_encodings[0]


def face_distance_to_confidence(face_distance: float, face_match_threshold: float) -> float:
    """
    Convert a face distance to a confidence score between 0.0 and 1.0.
    Based on the reference implementation from face_recognition docs.
    """
    if face_distance > face_match_threshold:
        range_val = 1.0 - face_match_threshold
        linear_val = (1.0 - face_distance) / (range_val * 2.0)
        return max(0.0, min(1.0, linear_val))

    range_val = face_match_threshold
    linear_val = 1.0 - (face_distance / (range_val * 2.0))
    return max(0.0, min(1.0, linear_val + ((1.0 - linear_val) * math.pow((linear_val - 0.5) * 2, 0.2))))


def compare_faces(stored_embedding: List[float], new_embedding: np.ndarray) -> tuple:
    """
    Compare two face embeddings and return match result and confidence.
    """
    stored_array = np.asarray(stored_embedding, dtype=np.float64)
    if stored_array.shape != (128,) or not np.isfinite(stored_array).all():
        raise HTTPException(
            status_code=400,
            detail="Invalid stored embedding values. Expected 128 numeric values."
        )
    
    # Calculate face distance (lower = more similar)
    face_distance = float(face_recognition.face_distance([stored_array], new_embedding)[0])
    
    # Convert distance to confidence (invert so higher = more confident match)
    confidence = face_distance_to_confidence(face_distance, FACE_MATCH_THRESHOLD)
    
    # Check if faces match based on threshold
    is_match = face_distance <= FACE_MATCH_THRESHOLD
    
    return is_match, float(confidence), face_distance


@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Health check endpoint."""
    return HealthResponse(
        status="healthy",
        version="1.0.0"
    )


@app.post("/enroll", response_model=EnrollResponse)
async def enroll_face(request: EnrollRequest):
    """
    Enroll a new face - extract and return face embedding for storage.
    
    The embedding should be stored in the database by the client app.
    """
    try:
        # Decode image
        image = decode_base64_image(request.image_base64)
        
        # Extract face encoding
        face_encoding = extract_face_encoding(image)
        
        # Generate a simple face ID (timestamp-based)
        import time
        face_id = f"face_{request.user_id}_{int(time.time() * 1000)}"
        
        # Convert encoding to list for JSON serialization
        embedding_list = face_encoding.tolist()
        
        return EnrollResponse(
            success=True,
            face_id=face_id,
            face_embedding=embedding_list,
            message="Face enrolled successfully. Store the embedding securely."
        )
    
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Face enrollment failed: {str(e)}"
        )


@app.post("/verify", response_model=VerifyResponse)
async def verify_face(request: VerifyRequest):
    """
    Verify a face against a stored embedding.
    
    Returns whether the faces match and the confidence score.
    """
    try:
        # Validate stored embedding
        if not request.stored_embedding or len(request.stored_embedding) != 128:
            raise HTTPException(
                status_code=400,
                detail="Invalid stored embedding. Expected 128-dimensional vector."
            )
        
        # Decode image
        image = decode_base64_image(request.image_base64)
        
        # Extract face encoding from new image
        face_encoding = extract_face_encoding(image)
        
        # Compare faces
        is_match, confidence, distance = compare_faces(request.stored_embedding, face_encoding)
        
        if is_match:
            message = f"Face verified successfully (confidence: {confidence:.1%})"
        else:
            message = f"Face verification failed - faces do not match (confidence: {confidence:.1%})"
        
        return VerifyResponse(
            success=is_match,
            match=is_match,
            confidence=confidence,
            distance=distance,
            message=message
        )
    
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Face verification failed: {str(e)}"
        )


@app.get("/")
async def root():
    """Root endpoint with API info."""
    return {
        "name": "Face Verification API",
        "version": "1.0.0",
        "endpoints": {
            "/health": "Health check",
            "/enroll": "Enroll a new face (POST)",
            "/verify": "Verify a face against stored embedding (POST)"
        }
    }


if __name__ == "__main__":
    import uvicorn
    port = int(os.getenv("PORT", "8000"))
    uvicorn.run(app, host="0.0.0.0", port=port)