"""
Face Recognition Service

Uses InsightFace with ONNX Runtime for:
- Face detection (SCRFD)
- Face embedding generation (ArcFace)
- Gender and age estimation

Provides REST API endpoints for the digest pipeline.
"""

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import insightface
import numpy as np
import base64
import cv2
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="Face Recognition Service",
    description="Face detection and embedding generation using InsightFace",
    version="1.0.0"
)

# Global face analyzer - loaded on startup
face_analyzer = None


@app.on_event("startup")
async def load_models():
    """Load InsightFace models on startup."""
    global face_analyzer
    logger.info("Loading InsightFace models...")

    try:
        # buffalo_l includes SCRFD (detection) + ArcFace (recognition)
        face_analyzer = insightface.app.FaceAnalysis(
            name='buffalo_l',
            providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
        )
        # Prepare with detection size
        face_analyzer.prepare(ctx_id=0, det_size=(640, 640))
        logger.info("InsightFace models loaded successfully")
    except Exception as e:
        logger.error(f"Failed to load InsightFace models: {e}")
        raise


class AnalyzeRequest(BaseModel):
    """Request body for face analysis."""
    image_base64: str


class Face(BaseModel):
    """Detected face with embedding and attributes."""
    bbox: list[float]  # [x1, y1, x2, y2]
    embedding: list[float]  # 512-D vector
    gender: int  # 0=female, 1=male
    age: int
    det_score: float  # Detection confidence


class AnalyzeResponse(BaseModel):
    """Response from face analysis."""
    faces: list[Face]
    count: int
    women_count: int
    men_count: int


class HealthResponse(BaseModel):
    """Health check response."""
    status: str
    model_loaded: bool


@app.post("/analyze", response_model=AnalyzeResponse)
async def analyze_faces(request: AnalyzeRequest):
    """
    Detect faces in an image and return embeddings with attributes.

    - **image_base64**: Base64 encoded image (JPEG or PNG)

    Returns detected faces with:
    - Bounding boxes
    - 512-D face embeddings (for recognition)
    - Gender (0=female, 1=male)
    - Age estimate
    - Detection confidence score
    """
    if face_analyzer is None:
        raise HTTPException(503, "Face analyzer not initialized")

    try:
        # Decode base64 image
        img_data = base64.b64decode(request.image_base64)
        img_array = np.frombuffer(img_data, np.uint8)
        img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)

        if img is None:
            raise HTTPException(400, "Invalid image data - could not decode")

        # Detect and analyze faces
        faces = face_analyzer.get(img)

        # Convert to response format
        face_results = []
        women_count = 0
        men_count = 0

        for face in faces:
            gender = int(face.gender) if hasattr(face, 'gender') else -1
            if gender == 0:
                women_count += 1
            elif gender == 1:
                men_count += 1

            face_results.append(Face(
                bbox=face.bbox.tolist(),
                embedding=face.embedding.tolist() if hasattr(face, 'embedding') else [],
                gender=gender,
                age=int(face.age) if hasattr(face, 'age') else -1,
                det_score=float(face.det_score) if hasattr(face, 'det_score') else 0.0
            ))

        logger.info(f"Detected {len(faces)} faces ({women_count} women, {men_count} men)")

        return AnalyzeResponse(
            faces=face_results,
            count=len(faces),
            women_count=women_count,
            men_count=men_count
        )

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Face analysis failed: {e}")
        raise HTTPException(500, f"Face analysis failed: {str(e)}")


@app.get("/health", response_model=HealthResponse)
async def health():
    """Health check endpoint."""
    return HealthResponse(
        status="ok",
        model_loaded=face_analyzer is not None
    )


@app.get("/")
async def root():
    """Root endpoint with service info."""
    return {
        "service": "Face Recognition",
        "version": "1.0.0",
        "endpoints": {
            "/analyze": "POST - Analyze faces in image",
            "/health": "GET - Health check"
        }
    }


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5001)