package com.example.studentfaceregistry.face

import android.content.Context
import android.graphics.Bitmap
import android.graphics.Rect
import com.google.mlkit.vision.common.InputImage
import com.google.mlkit.vision.face.Face
import com.google.mlkit.vision.face.FaceDetection
import com.google.mlkit.vision.face.FaceDetectorOptions
import com.google.mlkit.vision.face.FaceLandmark
import kotlinx.coroutines.tasks.await

/**
 * 人脸处理器
 * 负责人脸检测、对齐、预处理和特征提取
 */
class FaceProcessor(context: Context) : AutoCloseable {
    private val embedder = FaceEmbedder(context)
    private val aligner = FaceAligner()
    private val preprocessor = ImagePreprocessor()

    // 实时检测器：快速模式，用于摄像头预览
    private val realtimeDetector = FaceDetection.getClient(
        FaceDetectorOptions.Builder()
            .setPerformanceMode(FaceDetectorOptions.PERFORMANCE_MODE_FAST)
            .setMinFaceSize(0.06f)
            .enableTracking()
            .setLandmarkMode(FaceDetectorOptions.LANDMARK_MODE_ALL)
            .build()
    )

    // 注册检测器：高精度模式，启用关键点检测
    private val enrollmentDetector = FaceDetection.getClient(
        FaceDetectorOptions.Builder()
            .setPerformanceMode(FaceDetectorOptions.PERFORMANCE_MODE_ACCURATE)
            .setMinFaceSize(0.04f)
            .setLandmarkMode(FaceDetectorOptions.LANDMARK_MODE_ALL)
            .build()
    )

    /**
     * 从图片中提取单张人脸的特征（用于注册/上传）
     * @param bitmap 输入图片
     * @return L2 归一化的特征向量
     */
    suspend fun embedSingleFace(bitmap: Bitmap): FloatArray {
        val faces = detectFacesForEnrollment(bitmap)
            .mapNotNull { face ->
                val safeBox = safeCropRect(bitmap, face.boundingBox)
                if (safeBox != null) face to safeBox else null
            }

        require(faces.isNotEmpty()) { "图片中未检测到可用人脸，请换更清晰或更正面的照片。" }

        // 选择面积最大的人脸
        val (largestFace, faceBox) = faces.maxBy { (_, box) ->
            box.width() * box.height()
        }

        // 裁剪人脸
        val croppedFace = Bitmap.createBitmap(
            bitmap,
            faceBox.left,
            faceBox.top,
            faceBox.width(),
            faceBox.height()
        )

        // 使用关键点进行对齐
        val alignedFace = alignFaceWithLandmarks(croppedFace, largestFace, faceBox.left, faceBox.top)

        // 图像预处理
        val preprocessedFace = preprocessor.preprocess(alignedFace)

        // 提取特征
        return embedder.embed(preprocessedFace)
    }

    /**
     * 使用关键点对齐人脸
     */
    private fun alignFaceWithLandmarks(
        faceBitmap: Bitmap,
        face: Face,
        faceLeft: Int,
        faceTop: Int
    ): Bitmap {
        val leftEye = face.getLandmark(FaceLandmark.LEFT_EYE) ?: return faceBitmap
        val rightEye = face.getLandmark(FaceLandmark.RIGHT_EYE) ?: return faceBitmap

        return aligner.align(
            bitmap = faceBitmap,
            leftEyeX = leftEye.position.x,
            leftEyeY = leftEye.position.y,
            rightEyeX = rightEye.position.x,
            rightEyeY = rightEye.position.y,
            faceLeft = faceLeft,
            faceTop = faceTop
        )
    }

    /**
     * 实时检测人脸（用于摄像头预览）
     */
    suspend fun detectFaces(image: InputImage): List<Face> {
        return realtimeDetector.process(image).await()
    }

    /**
     * 实时检测人脸（bitmap 版本）
     */
    suspend fun detectFaces(bitmap: Bitmap): List<Face> {
        return detectFaces(InputImage.fromBitmap(bitmap, 0))
    }

    /**
     * 注册用的人脸检测（高精度）
     */
    private suspend fun detectFacesForEnrollment(bitmap: Bitmap): List<Face> {
        val image = InputImage.fromBitmap(bitmap, 0)
        val accurateFaces = enrollmentDetector.process(image).await()
        if (accurateFaces.isNotEmpty()) return accurateFaces
        return realtimeDetector.process(image).await()
    }

    /**
     * 提取指定人脸的特征
     */
    fun embedFace(bitmap: Bitmap, face: Face): FloatArray {
        val crop = requireNotNull(safeCropRect(bitmap, face.boundingBox)) { "人脸区域超出图片边界。" }
        val cropped = Bitmap.createBitmap(bitmap, crop.left, crop.top, crop.width(), crop.height())
        val aligned = alignFaceForRecognition(cropped, face, crop.left, crop.top)
        val preprocessed = preprocessor.preprocessRealtime(aligned)
        return embedder.embed(preprocessed)
    }

    fun embedFaceCandidates(bitmap: Bitmap, face: Face): List<FloatArray> {
        val embeddings = mutableListOf<FloatArray>()
        for (paddingRatio in FACE_PADDING_RATIOS) {
            val crop = safeCropRect(bitmap, face.boundingBox, paddingRatio) ?: continue
            val cropped = Bitmap.createBitmap(bitmap, crop.left, crop.top, crop.width(), crop.height())
            val aligned = alignFaceForRecognition(cropped, face, crop.left, crop.top)
            val preprocessed = preprocessor.preprocessRealtime(aligned)
            embeddings.add(embedder.embed(preprocessed))
        }
        return embeddings
    }

    fun warmUp() {
        embedder.warmUp()
    }

    override fun close() {
        realtimeDetector.close()
        enrollmentDetector.close()
        embedder.close()
    }

    /**
     * 计算安全的裁剪区域（带 padding）
     */
    private fun safeCropRect(bitmap: Bitmap, box: Rect): Rect? {
        return safeCropRect(bitmap, box, 0.25f)
    }

    private fun safeCropRect(bitmap: Bitmap, box: Rect, paddingRatio: Float): Rect? {
        if (box.width() <= 0 || box.height() <= 0) return null

        // 增加 padding 比例，保留更多上下文信息
        val padding = (maxOf(box.width(), box.height()) * paddingRatio).toInt()

        val left = (box.left - padding).coerceIn(0, bitmap.width)
        val top = (box.top - padding).coerceIn(0, bitmap.height)
        val right = (box.right + padding).coerceIn(0, bitmap.width)
        val bottom = (box.bottom + padding).coerceIn(0, bitmap.height)

        val width = right - left
        val height = bottom - top

        if (width <= 1 || height <= 1) return null

        return Rect(left, top, right, bottom)
    }

    private fun alignFaceForRecognition(
        faceBitmap: Bitmap,
        face: Face,
        faceLeft: Int,
        faceTop: Int
    ): Bitmap {
        val leftEye = face.getLandmark(FaceLandmark.LEFT_EYE)
        val rightEye = face.getLandmark(FaceLandmark.RIGHT_EYE)
        if (leftEye != null && rightEye != null) {
            return aligner.align(
                bitmap = faceBitmap,
                leftEyeX = leftEye.position.x,
                leftEyeY = leftEye.position.y,
                rightEyeX = rightEye.position.x,
                rightEyeY = rightEye.position.y,
                faceLeft = faceLeft,
                faceTop = faceTop
            )
        }
        return aligner.alignWithEstimation(faceBitmap, 0, 0, face.headEulerAngleZ)
    }

    companion object {
        private val FACE_PADDING_RATIOS = floatArrayOf(0.18f, 0.34f)
    }
}
