package com.example.studentfaceregistry.face

import android.content.Context
import android.graphics.Bitmap
import android.graphics.BitmapFactory
import android.graphics.Rect
import android.media.MediaDataSource
import android.media.MediaMetadataRetriever
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
import kotlin.math.PI
import kotlin.math.atan2
import kotlin.math.sqrt

/**
 * 智能注册器
 * 支持图片和视频两种输入方式，自动选择最大人脸进行注册
 *
 * 功能特性：
 * 1. 自动识别输入类型（图片或视频）
 * 2. 自动检测并选择最大的人脸
 * 3. 视频注册时自动提取多角度帧并融合特征
 */
class SmartEnrollment(context: Context) : AutoCloseable {

    // 高精度人脸检测器（用于注册）
    private val detector = FaceDetection.getClient(
        FaceDetectorOptions.Builder()
            .setPerformanceMode(FaceDetectorOptions.PERFORMANCE_MODE_ACCURATE)
            .setLandmarkMode(FaceDetectorOptions.LANDMARK_MODE_ALL)
            .setMinFaceSize(0.03f)
            .build()
    )

    private val embedder = FaceEmbedder(context, FaceEmbedder.ModelType.ARCFACE)
    private val aligner = FaceAligner()
    private val preprocessor = ImagePreprocessor()

    /**
     * 智能注册入口
     * @param data 数据字节（图片或视频）
     * @param mimeType MIME 类型，如 "image/jpeg", "video/mp4"
     * @param fileName 文件名（用于提取学号和姓名）
     * @return 注册结果
     */
    suspend fun enroll(
        data: ByteArray,
        mimeType: String?,
        fileName: String
    ): EnrollmentResult {
        return if (isVideo(mimeType, fileName)) {
            enrollFromVideo(data)
        } else {
            enrollFromImage(data)
        }
    }

    /**
     * 从图片注册
     */
    private suspend fun enrollFromImage(data: ByteArray): EnrollmentResult {
        return try {
            // 解码图片
            val bitmap = decodeBitmap(data)

            // 检测人脸并只保留可安全裁剪的人脸区域
            val faces = detectUsableFaces(bitmap)

            if (faces.isEmpty()) {
                return EnrollmentResult.Failure("未检测到人脸，请换更清晰或更正面的照片。")
            }

            // 选择最大的人脸
            val largestFace = faces.maxBy { faceInfo ->
                faceInfo.faceArea
            }

            val embeddings = mutableListOf<FloatArray>()
            for (paddingRatio in IMAGE_PADDING_RATIOS) {
                val cropRect = safeCropRect(bitmap, largestFace.face.boundingBox, paddingRatio) ?: continue
                val croppedFace = cropFace(bitmap, cropRect)
                val alignedFace = alignFace(croppedFace, largestFace.face, cropRect)
                val preprocessedFace = preprocessor.preprocess(alignedFace)
                embeddings.add(embedder.embed(preprocessedFace))
            }

            if (embeddings.isEmpty()) {
                return EnrollmentResult.Failure("未检测到可用人脸，请换更清晰或更正面的照片。")
            }

            val embedding = fuseEmbeddings(embeddings)

            EnrollmentResult.Success(embedding, EnrollmentSourceType.IMAGE)
        } catch (e: Exception) {
            EnrollmentResult.Failure("图片处理失败：${e.message}")
        }
    }

    /**
     * 从视频注册
     * @param data 视频字节
     * @param minFrames 最少帧数
     * @param maxFrames 最大帧数（用于融合）
     */
    private suspend fun enrollFromVideo(
        data: ByteArray,
        minFrames: Int = 10,
        maxFrames: Int = 30
    ): EnrollmentResult {
        return try {
            // 提取视频帧
            val frames = extractVideoFrames(data, maxFrames * 2)

            if (frames.size < minFrames) {
                return EnrollmentResult.Failure(
                    "视频有效帧数不足（检测到 ${frames.size} 帧，需要至少 $minFrames 帧），" +
                    "请确保视频中人脸清晰且持续展示足够时间。"
                )
            }

            // 对每帧选择最大人脸
            val faceFrames = frames.mapNotNull { bitmap ->
                val faces = detectUsableFaces(bitmap)
                if (faces.isEmpty()) return@mapNotNull null

                val largestFace = faces.maxBy { faceInfo ->
                    faceInfo.faceArea
                }

                FaceFrameInfo(bitmap, largestFace.face, largestFace.cropRect, largestFace.faceArea)
            }

            if (faceFrames.size < minFrames) {
                return EnrollmentResult.Failure(
                    "视频有效人脸帧数不足（检测到 ${faceFrames.size} 帧，需要至少 $minFrames 帧）"
                )
            }

            // 计算每帧的角度并分组
            val groupedFrames = groupFramesByAngle(faceFrames)

            // 从每组选择质量最好的帧
            val selectedFrames = selectBestFramesFromGroups(groupedFrames, maxFrames)

            // 提取每帧的特征
            val embeddings = selectedFrames.map { frameInfo ->
                val cropped = cropFace(frameInfo.bitmap, frameInfo.cropRect)
                val aligned = alignFace(cropped, frameInfo.face, frameInfo.cropRect)
                val preprocessed = preprocessor.preprocess(aligned)
                embedder.embed(preprocessed)
            }

            // 融合特征
            val fusedEmbedding = fuseEmbeddings(embeddings)

            // 计算角度覆盖
            val yawRange = calculateAngleRange(selectedFrames, { f -> yawAngle(f.face) })
            val pitchRange = calculateAngleRange(selectedFrames, { f -> pitchAngle(f.face) })

            EnrollmentResult.Success(
                embedding = fusedEmbedding,
                sourceType = EnrollmentSourceType.VIDEO,
                frameCount = selectedFrames.size,
                yawRange = yawRange,
                pitchRange = pitchRange
            )
        } catch (e: Exception) {
            EnrollmentResult.Failure("视频处理失败：${e.message}")
        }
    }

    /**
     * 提取视频帧
     */
    private fun extractVideoFrames(videoData: ByteArray, maxFrames: Int): List<Bitmap> {
        val retriever = MediaMetadataRetriever()
        val dataSource = ByteArrayVideoDataSource(videoData)

        return try {
            retriever.setDataSource(dataSource)

            val duration = retriever.extractMetadata(MediaMetadataRetriever.METADATA_KEY_DURATION)?.toLongOrNull() ?: 0
            val frameInterval = maxOf(100, (duration / maxFrames).toInt())  // 至少每 100ms 一帧

            val frames = mutableListOf<Bitmap>()
            var timeUs = 0L

            while (timeUs < duration * 1000) {
                try {
                    val bitmap = retriever.getFrameAtTime(timeUs, MediaMetadataRetriever.OPTION_CLOSEST)
                    if (bitmap != null && !bitmap.isRecycled) {
                        frames.add(bitmap)
                    }
                } catch (e: Exception) {
                    // 跳过错误帧
                }
                timeUs += (frameInterval * 1000).toLong()  // 转换为微秒
            }

            frames
        } finally {
            retriever.release()
            dataSource.close()
        }
    }

    /**
     * 计算偏航角（左右转头）
     */
    private fun yawAngle(face: Face): Float {
        val leftCheek = face.getLandmark(FaceLandmark.LEFT_CHEEK) ?: return 0f
        val rightCheek = face.getLandmark(FaceLandmark.RIGHT_CHEEK) ?: return 0f

        val dx = rightCheek.position.x - leftCheek.position.x
        val dy = rightCheek.position.y - leftCheek.position.y
        return atan2(dy.toDouble(), dx.toDouble()).toFloat() * 180f / PI.toFloat()
    }

    /**
     * 计算俯仰角（上下点头）
     */
    private fun pitchAngle(face: Face): Float {
        val nose = face.getLandmark(FaceLandmark.NOSE_BASE) ?: return 0f
        val leftEye = face.getLandmark(FaceLandmark.LEFT_EYE) ?: return 0f
        val rightEye = face.getLandmark(FaceLandmark.RIGHT_EYE) ?: return 0f

        val eyeCenterY = (leftEye.position.y + rightEye.position.y) / 2f
        val dx = nose.position.x - (leftEye.position.x + rightEye.position.x) / 2f
        val dy = nose.position.y - eyeCenterY
        return atan2(dy.toDouble(), dx.toDouble()).toFloat() * 180f / PI.toFloat()
    }

    /**
     * 按角度分组
     */
    private fun groupFramesByAngle(frames: List<FaceFrameInfo>): Map<AngleBucket, List<FaceFrameInfo>> {
        val buckets = mutableMapOf<AngleBucket, MutableList<FaceFrameInfo>>()

        for (frame in frames) {
            val yaw = yawAngle(frame.face)
            val pitch = pitchAngle(frame.face)
            val bucket = AngleBucket(
                yawBucket = when {
                    yaw < -30 -> -2
                    yaw < -10 -> -1
                    yaw < 10 -> 0
                    yaw < 30 -> 1
                    else -> 2
                },
                pitchBucket = when {
                    pitch < -15 -> -1
                    pitch < 15 -> 0
                    else -> 1
                }
            )
            if (!buckets.containsKey(bucket)) buckets[bucket] = mutableListOf()
            buckets[bucket]!!.add(frame)
        }

        return buckets
    }

    /**
     * 从每组选择最佳帧
     */
    private fun selectBestFramesFromGroups(
        groups: Map<AngleBucket, List<FaceFrameInfo>>,
        maxFrames: Int
    ): List<FaceFrameInfo> {
        val selected = mutableListOf<FaceFrameInfo>()

        // 按组内最大人脸面积排序
        val sortedGroups = groups.entries.sortedByDescending { entry ->
            entry.value.maxOfOrNull {
                it.faceArea
            } ?: 0
        }

        for ((_, frames) in sortedGroups) {
            val sortedFrames = frames.sortedByDescending {
                it.faceArea
            }
            for (frame in sortedFrames) {
                if (selected.size >= maxFrames) break
                selected.add(frame)
            }
        }

        return selected
    }

    /**
     * 融合多个特征向量
     */
    private fun fuseEmbeddings(embeddings: List<FloatArray>): FloatArray {
        if (embeddings.isEmpty()) throw IllegalArgumentException("embeddings is empty")
        if (embeddings.size == 1) return embeddings[0]

        val size = embeddings[0].size
        val sum = FloatArray(size)

        for (embedding in embeddings) {
            for (i in 0 until size) {
                sum[i] += embedding[i]
            }
        }

        // 平均
        for (i in 0 until size) {
            sum[i] /= embeddings.size
        }

        // L2 归一化
        return l2Normalize(sum)
    }

    /**
     * 计算角度范围
     */
    private fun calculateAngleRange(
        frames: List<FaceFrameInfo>,
        angleExtractor: (FaceFrameInfo) -> Float
    ): Float {
        if (frames.isEmpty()) return 0f
        val angles = frames.map(angleExtractor)
        return angles.maxOrNull()!! - angles.minOrNull()!!
    }

    /**
     * 裁剪人脸
     */
    private fun cropFace(bitmap: Bitmap, box: Rect): Bitmap {
        return Bitmap.createBitmap(bitmap, box.left, box.top, box.width(), box.height())
    }

    private suspend fun detectUsableFaces(bitmap: Bitmap): List<DetectedFace> {
        return detector.process(InputImage.fromBitmap(bitmap, 0)).await()
            .mapNotNull { face ->
                val cropRect = safeCropRect(bitmap, face.boundingBox) ?: return@mapNotNull null
                val faceArea = face.boundingBox.width() * face.boundingBox.height()
                DetectedFace(face, cropRect, faceArea)
            }
    }

    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

        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)
        if (right - left <= 1 || bottom - top <= 1) return null

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

    /**
     * 对齐人脸
     */
    private fun alignFace(bitmap: Bitmap, face: Face, cropRect: Rect): Bitmap {
        val leftEye = face.getLandmark(FaceLandmark.LEFT_EYE) ?: return bitmap
        val rightEye = face.getLandmark(FaceLandmark.RIGHT_EYE) ?: return bitmap

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

    /**
     * 解码 Bitmap
     */
    private fun decodeBitmap(data: ByteArray): Bitmap {
        val bounds = BitmapFactory.Options().apply { inJustDecodeBounds = true }
        BitmapFactory.decodeByteArray(data, 0, data.size, bounds)

        val maxSide = maxOf(bounds.outWidth, bounds.outHeight).coerceAtLeast(1)
        var sampleSize = 1
        while (maxSide / sampleSize > 1280) sampleSize *= 2

        val options = BitmapFactory.Options().apply {
            inSampleSize = sampleSize
            inPreferredConfig = Bitmap.Config.ARGB_8888
        }
        return BitmapFactory.decodeByteArray(data, 0, data.size, options)
            ?: throw IllegalArgumentException("Unable to decode image")
    }

    /**
     * L2 归一化
     */
    private fun l2Normalize(values: FloatArray): FloatArray {
        var sum = 0f
        for (value in values) sum += value * value
        val norm = sqrt(sum.coerceAtLeast(1e-12f))
        return FloatArray(values.size) { values[it] / norm }
    }

    /**
     * 判断是否为视频
     */
    private fun isVideo(mimeType: String?, fileName: String): Boolean {
        val lowerMimeType = mimeType?.lowercase().orEmpty()
        val lowerFileName = fileName.lowercase()
        return lowerMimeType.startsWith("video/") ||
               lowerMimeType in VIDEO_MIME_TYPES ||
               VIDEO_EXTENSIONS.any { lowerFileName.endsWith(it) }
    }

    override fun close() {
        detector.close()
        embedder.close()
    }

    companion object {
        private val VIDEO_EXTENSIONS = setOf(".mp4", ".avi", ".mov", ".wmv", ".flv", ".mkv", ".webm", ".3gp")
        private val VIDEO_MIME_TYPES = setOf("application/mp4")
        private val IMAGE_PADDING_RATIOS = floatArrayOf(0.18f, 0.25f, 0.34f)
    }
}

private data class DetectedFace(
    val face: Face,
    val cropRect: Rect,
    val faceArea: Int
)

/**
 * 人脸帧信息
 */
private data class FaceFrameInfo(
    val bitmap: Bitmap,
    val face: Face,
    val cropRect: Rect,
    val faceArea: Int
)

private class ByteArrayVideoDataSource(private val data: ByteArray) : MediaDataSource() {
    override fun readAt(position: Long, buffer: ByteArray, offset: Int, size: Int): Int {
        if (position >= data.size) return -1
        val length = minOf(size, data.size - position.toInt())
        System.arraycopy(data, position.toInt(), buffer, offset, length)
        return length
    }

    override fun getSize(): Long = data.size.toLong()

    override fun close() = Unit
}

/**
 * 角度分组
 */
private data class AngleBucket(val yawBucket: Int, val pitchBucket: Int)

/**
 * 注册结果
 */
sealed class EnrollmentResult {
    data class Success(
        val embedding: FloatArray,
        val sourceType: EnrollmentSourceType = EnrollmentSourceType.IMAGE,
        val frameCount: Int = 1,
        val yawRange: Float = 0f,
        val pitchRange: Float = 0f
    ) : EnrollmentResult() {
        override fun equals(other: Any?): Boolean {
            if (this === other) return true
            if (javaClass != other?.javaClass) return false

            other as Success

            if (frameCount != other.frameCount) return false
            if (yawRange != other.yawRange) return false
            if (pitchRange != other.pitchRange) return false
            if (!embedding.contentEquals(other.embedding)) return false
            if (sourceType != other.sourceType) return false

            return true
        }

        override fun hashCode(): Int {
            var result = frameCount
            result = 31 * result + yawRange.hashCode()
            result = 31 * result + pitchRange.hashCode()
            result = 31 * result + embedding.contentHashCode()
            result = 31 * result + sourceType.hashCode()
            return result
        }
    }

    data class Failure(val reason: String) : EnrollmentResult()
}

/**
 * 注册源类型
 */
enum class EnrollmentSourceType {
    IMAGE,   // 图片注册
    VIDEO    // 视频注册
}
