"""
DeepFace Service — Facial analysis for video frames.

Model: DeepFace (opencv backend)
Endpoint: POST /analyze
"""

import base64
import io
import logging
from contextlib import asynccontextmanager
from typing import Optional

import cv2
import numpy as np
from fastapi import FastAPI
from PIL import Image
from pydantic import BaseModel

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


@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    DeepFace downloads model weights lazily on first call, so we do a
    warm-up inference during startup to avoid a cold-start delay on the
    first real request.
    """
    logger.info("Warming up DeepFace models …")
    try:
        from deepface import DeepFace  # noqa: PLC0415

        # 1×1 black image — just enough to trigger weight download without
        # actually detecting a face (the exception is silently ignored).
        dummy = np.zeros((64, 64, 3), dtype=np.uint8)
        try:
            DeepFace.analyze(dummy, actions=["emotion", "age", "gender"], enforce_detection=False, detector_backend="retinaface")
        except Exception:  # noqa: BLE001
            pass
        logger.info("DeepFace warm-up done")
    except Exception as exc:  # noqa: BLE001
        logger.warning("DeepFace warm-up failed (non-fatal): %s", exc)

    yield


app = FastAPI(title="ViralClip DeepFace Service", lifespan=lifespan)


# ---------------------------------------------------------------------------
# Request / Response schemas
# ---------------------------------------------------------------------------

class BBox(BaseModel):
    x: int
    y: int
    w: int
    h: int


class EmotionResult(BaseModel):
    dominant: str
    scores: dict[str, float]


class FaceResult(BaseModel):
    face_index: int
    emotion: EmotionResult
    age: int
    gender: str
    gender_confidence: float
    bbox: BBox


class AnalyzeRequest(BaseModel):
    frame_b64: str


class AnalyzeResponse(BaseModel):
    faces: list[FaceResult]
    face_count: int
    dominant_emotion: Optional[str]
    frame_width: Optional[int] = None
    frame_height: Optional[int] = None


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _parse_face(idx: int, face_data: dict) -> FaceResult:
    """Convert a single DeepFace result dict into a FaceResult."""
    # Emotion scores
    raw_emotions: dict = face_data.get("emotion", {})
    emotion_scores = {k: round(float(v) / 100.0, 4) for k, v in raw_emotions.items()}
    dominant_emotion: str = face_data.get("dominant_emotion", max(emotion_scores, key=lambda k: emotion_scores[k]))

    # Gender
    gender_raw: dict = face_data.get("gender", {})
    if isinstance(gender_raw, dict):
        # {"Man": 94.2, "Woman": 5.8}
        dominant_gender = max(gender_raw, key=lambda k: gender_raw[k])
        gender_conf = round(float(gender_raw[dominant_gender]) / 100.0, 4)
    else:
        dominant_gender = str(gender_raw)
        gender_conf = float(face_data.get("gender_confidence", 1.0))

    # BBox
    region: dict = face_data.get("region", {})
    bbox = BBox(
        x=int(region.get("x", 0)),
        y=int(region.get("y", 0)),
        w=int(region.get("w", 0)),
        h=int(region.get("h", 0)),
    )

    return FaceResult(
        face_index=idx,
        emotion=EmotionResult(dominant=dominant_emotion, scores=emotion_scores),
        age=int(face_data.get("age", 0)),
        gender=dominant_gender,
        gender_confidence=gender_conf,
        bbox=bbox,
    )


# ---------------------------------------------------------------------------
# Endpoint
# ---------------------------------------------------------------------------

@app.post("/analyze")
async def analyze(req: AnalyzeRequest):
    frame_width: Optional[int] = None
    frame_height: Optional[int] = None
    try:
        from deepface import DeepFace  # noqa: PLC0415

        # Decode image → numpy BGR array — pad base64 if needed
        b64 = req.frame_b64
        b64 += "=" * (-len(b64) % 4)
        img_bytes = base64.b64decode(b64)
        pil_image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
        frame_width, frame_height = pil_image.size
        frame_rgb = np.array(pil_image)
        frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)

        results = DeepFace.analyze(
            frame_bgr,
            actions=["emotion", "age", "gender"],
            enforce_detection=True,
            detector_backend="retinaface",
        )

        # DeepFace may return a single dict or a list of dicts
        if isinstance(results, dict):
            results = [results]

        faces = [_parse_face(i, r) for i, r in enumerate(results)]

        # dominant_emotion: from the face whose top emotion score is highest
        dominant_emotion: Optional[str] = None
        if faces:
            best_face = max(
                faces,
                key=lambda f: f.emotion.scores.get(f.emotion.dominant, 0.0),
            )
            dominant_emotion = best_face.emotion.dominant

        return AnalyzeResponse(
            faces=faces,
            face_count=len(faces),
            dominant_emotion=dominant_emotion,
            frame_width=frame_width,
            frame_height=frame_height,
        )

    except Exception as exc:  # noqa: BLE001
        err_msg = str(exc).lower()
        # DeepFace raises ValueError or a custom exception when no face is found
        if "face" in err_msg and ("not" in err_msg or "detect" in err_msg or "found" in err_msg):
            # frame_width/frame_height are set above the DeepFace.analyze
            # call (which is what raises this), so they're populated
            # unless image decoding itself failed — in which case they
            # stay at their None default.
            return AnalyzeResponse(
                faces=[], face_count=0, dominant_emotion=None,
                frame_width=frame_width, frame_height=frame_height,
            )

        logger.exception("DeepFace inference failed")
        return {"error": str(exc)}


@app.get("/health")
async def health():
    return {"status": "ok", "model": "deepface"}