app/models/schemas.py

"""
Pydantic schemas for KYC POC API.

This module defines all request and response models for the API.
"""

from pydantic import BaseModel, Field
from typing import Optional, Dict, Any


# ============================================================================
# Common Models
# ============================================================================

class BoundingBox(BaseModel):
    """Face bounding box coordinates."""
    x: int = Field(..., description="X coordinate of top-left corner")
    y: int = Field(..., description="Y coordinate of top-left corner")
    width: int = Field(..., description="Width of bounding box")
    height: int = Field(..., description="Height of bounding box")


class FacePose(BaseModel):
    """Face pose angles."""
    yaw: float = Field(..., description="Yaw angle (left-right rotation)")
    pitch: float = Field(..., description="Pitch angle (up-down rotation)")
    roll: float = Field(..., description="Roll angle (tilt)")
    is_frontal: bool = Field(..., description="Whether face is frontal")


class Demographics(BaseModel):
    """Demographic information extracted from face."""
    age: Optional[int] = Field(None, description="Estimated age")
    gender: Optional[str] = Field(None, description="Estimated gender (Male/Female)")


# ============================================================================
# Quality Models
# ============================================================================

class BlurAnalysis(BaseModel):
    """Blur analysis results."""
    blur_score: float = Field(..., description="Laplacian variance score (higher = sharper)")
    blur_threshold: float = Field(..., description="Threshold below which image is considered blurry")
    is_blurry: bool = Field(..., description="Whether image is too blurry")


class BrightnessAnalysis(BaseModel):
    """Brightness analysis results."""
    brightness: float = Field(..., description="Mean brightness (0-1)")
    brightness_min: float = Field(..., description="Minimum acceptable brightness")
    brightness_max: float = Field(..., description="Maximum acceptable brightness")
    is_too_dark: bool = Field(..., description="Whether image is too dark")
    is_too_bright: bool = Field(..., description="Whether image is too bright")


class QualityAnalysis(BaseModel):
    """Complete quality analysis for an image."""
    blur_score: float
    blur_threshold: float
    is_blurry: bool
    brightness: float
    brightness_min: float
    brightness_max: float
    is_too_dark: bool
    is_too_bright: bool
    pose: Optional[FacePose] = None
    is_good_quality: bool = Field(..., description="Overall quality assessment")


# ============================================================================
# Face Match Models
# ============================================================================

class FaceMatchResult(BaseModel):
    """Face matching result."""
    is_match: bool = Field(..., description="Whether faces match")
    similarity_score: float = Field(..., description="Similarity score (0-1)")
    threshold: float = Field(..., description="Threshold used for matching")


# ============================================================================
# Liveness Models
# ============================================================================

class LivenessResult(BaseModel):
    """Liveness detection result."""
    is_real: bool = Field(..., description="Whether face is from a real person")
    confidence: float = Field(..., description="Confidence score")
    label: str = Field(..., description="Classification label")
    prediction_class: Optional[int] = Field(None, description="Raw prediction class")
    models_used: Optional[int] = Field(None, description="Number of models used")
    error: Optional[str] = Field(None, description="Error message if any")


# ============================================================================
# Request Models (Base64)
# ============================================================================

class Base64VerifyRequest(BaseModel):
    """Request model for base64 verification endpoint."""
    ktp_image: str = Field(..., description="Base64 encoded KTP image")
    selfie_image: str = Field(..., description="Base64 encoded selfie image")
    threshold: float = Field(default=0.5, ge=0.0, le=1.0, description="Similarity threshold")


class Base64FaceMatchRequest(BaseModel):
    """Request model for base64 face match endpoint."""
    image1: str = Field(..., description="Base64 encoded first image")
    image2: str = Field(..., description="Base64 encoded second image")
    threshold: float = Field(default=0.5, ge=0.0, le=1.0, description="Similarity threshold")


class Base64SingleImageRequest(BaseModel):
    """Request model for single image base64 endpoints."""
    image: str = Field(..., description="Base64 encoded image")


# ============================================================================
# Response Models
# ============================================================================

class ErrorResponse(BaseModel):
    """Error response model."""
    error_code: str = Field(..., description="Error code")
    message: str = Field(..., description="Error message")
    detail: Optional[str] = Field(None, description="Additional details")


class FaceInfo(BaseModel):
    """Information about a detected face."""
    bbox: BoundingBox
    demographics: Optional[Demographics] = None
    det_score: Optional[float] = Field(None, description="Detection confidence")


class QualityResponse(BaseModel):
    """Response for quality check endpoint."""
    success: bool
    quality: QualityAnalysis
    face_box: Optional[BoundingBox] = None
    demographics: Optional[Demographics] = None
    message: str


class LivenessResponse(BaseModel):
    """Response for liveness check endpoint."""
    success: bool
    liveness: LivenessResult
    message: str


class FaceMatchResponse(BaseModel):
    """Response for face match endpoint."""
    success: bool
    face_match: FaceMatchResult
    face1: Optional[FaceInfo] = None
    face2: Optional[FaceInfo] = None
    message: str


class VerifyResponse(BaseModel):
    """Response for full KYC verification endpoint."""
    success: bool
    face_match: FaceMatchResult
    liveness: LivenessResult
    quality: Dict[str, QualityAnalysis] = Field(
        ...,
        description="Quality analysis for each image (ktp, selfie)"
    )
    demographics: Dict[str, Demographics] = Field(
        ...,
        description="Demographics for each image (ktp, selfie)"
    )
    face_boxes: Dict[str, BoundingBox] = Field(
        ...,
        description="Face bounding boxes for each image (ktp, selfie)"
    )
    message: str


class HealthResponse(BaseModel):
    """Health check response."""
    status: str = Field(..., description="Service status")
    models_loaded: Dict[str, bool] = Field(
        ...,
        description="Status of each ML model"
    )
    version: str = Field(..., description="API version")


# ============================================================================
# OCR Models
# ============================================================================

class OCRFieldResult(BaseModel):
    """Result for a single OCR field."""
    value: str = Field(..., description="Extracted and sanitized value")
    confidence: float = Field(..., description="OCR confidence score (0-1)")
    raw_value: str = Field(..., description="Raw extracted text before sanitization")


class NIKValidation(BaseModel):
    """NIK validation result with extracted information."""
    is_valid: bool = Field(..., description="Whether NIK passes validation")
    errors: list = Field(default_factory=list, description="List of validation errors")
    extracted: Dict[str, Any] = Field(
        default_factory=dict,
        description="Information extracted from NIK (province code, city code, birth date, gender)"
    )


class KTPOCRData(BaseModel):
    """Structured KTP data extracted from OCR."""
    provinsi: Optional[OCRFieldResult] = Field(None, description="Province name")
    kabupaten_kota: Optional[OCRFieldResult] = Field(None, description="City/Regency name")
    nik: Optional[OCRFieldResult] = Field(None, description="NIK (16-digit ID number)")
    nama: Optional[OCRFieldResult] = Field(None, description="Full name")
    tempat_lahir: Optional[OCRFieldResult] = Field(None, description="Place of birth")
    tanggal_lahir: Optional[OCRFieldResult] = Field(None, description="Date of birth (DD-MM-YYYY)")
    jenis_kelamin: Optional[OCRFieldResult] = Field(None, description="Gender (LAKI-LAKI/PEREMPUAN)")
    golongan_darah: Optional[OCRFieldResult] = Field(None, description="Blood type")
    alamat: Optional[OCRFieldResult] = Field(None, description="Address")
    rt_rw: Optional[OCRFieldResult] = Field(None, description="RT/RW")
    kelurahan_desa: Optional[OCRFieldResult] = Field(None, description="Village/District")
    kecamatan: Optional[OCRFieldResult] = Field(None, description="Sub-district")
    agama: Optional[OCRFieldResult] = Field(None, description="Religion")
    status_perkawinan: Optional[OCRFieldResult] = Field(None, description="Marital status")
    pekerjaan: Optional[OCRFieldResult] = Field(None, description="Occupation")
    kewarganegaraan: Optional[OCRFieldResult] = Field(None, description="Nationality (WNI/WNA)")
    berlaku_hingga: Optional[OCRFieldResult] = Field(None, description="Valid until (date or SEUMUR HIDUP)")


class OCRTextBlock(BaseModel):
    """Raw OCR text block with position."""
    text: str = Field(..., description="Extracted text")
    confidence: float = Field(..., description="OCR confidence score")
    bbox: list = Field(..., description="Bounding box coordinates [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]")


class KTPValidation(BaseModel):
    """KTP data validation results."""
    nik: Optional[NIKValidation] = Field(None, description="NIK validation result")


class Base64OCRRequest(BaseModel):
    """Request model for base64 OCR endpoint."""
    image: str = Field(..., description="Base64 encoded KTP image")
    validate: bool = Field(default=True, description="Whether to validate extracted data (e.g., NIK)")


class OCRResponse(BaseModel):
    """Response for KTP OCR endpoint."""
    success: bool = Field(..., description="Whether OCR extraction was successful")
    data: KTPOCRData = Field(..., description="Structured KTP data")
    raw_text: list[OCRTextBlock] = Field(
        default_factory=list,
        description="Raw OCR results with bounding boxes"
    )
    validation: Optional[KTPValidation] = Field(None, description="Data validation results")
    message: str = Field(..., description="Result message")