package com.placeholder.sherpai2.data.repository

import android.content.Context
import android.graphics.Bitmap
import android.util.Log
import com.placeholder.sherpai2.data.local.dao.FaceModelDao
import com.placeholder.sherpai2.data.local.dao.ImageDao
import com.placeholder.sherpai2.data.local.dao.PersonAgeTagDao
import com.placeholder.sherpai2.data.local.dao.PersonDao
import com.placeholder.sherpai2.data.local.dao.PhotoFaceTagDao
import com.placeholder.sherpai2.data.local.entity.*
import com.placeholder.sherpai2.ml.FaceNetModel
import com.placeholder.sherpai2.ui.trainingprep.TrainingSanityChecker
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.map
import kotlinx.coroutines.withContext
import javax.inject.Inject
import javax.inject.Singleton

/**
 * FaceRecognitionRepository - Complete face recognition system
 *
 * USES STRING IDs TO MATCH YOUR SCHEMA:
 * - PersonEntity.id: String (UUID)
 * - ImageEntity.imageId: String
 * - FaceModelEntity.id: String (UUID)
 * - PhotoFaceTagEntity.id: String (UUID)
 */
@Singleton
class FaceRecognitionRepository @Inject constructor(
    private val context: Context,
    private val personDao: PersonDao,
    private val imageDao: ImageDao,
    private val faceModelDao: FaceModelDao,
    private val photoFaceTagDao: PhotoFaceTagDao,
    private val personAgeTagDao: PersonAgeTagDao
) {
    companion object {
        private const val TAG = "FaceRecognitionRepo"
    }

    private val faceNetModel by lazy { FaceNetModel(context) }

    // ======================
    // TRAINING OPERATIONS
    // ======================

    /**
     * Create a new person with face model in one operation.
     *
     * @return PersonId (String UUID)
     */
    suspend fun createPersonWithFaceModel(
        personName: String,
        validImages: List<TrainingSanityChecker.ValidTrainingImage>,
        onProgress: (Int, Int) -> Unit = { _, _ -> }
    ): String = withContext(Dispatchers.IO) {

        // Create PersonEntity with UUID
        val person = PersonEntity.create(name = personName)

        personDao.insert(person)

        // Train face model
        trainPerson(
            personId = person.id,
            validImages = validImages,
            onProgress = onProgress
        )

        person.id
    }

    /**
     * Train a face recognition model for an existing person.
     *
     * @param personId String UUID
     * @return Face model ID (String UUID)
     */
    suspend fun trainPerson(
        personId: String,
        validImages: List<TrainingSanityChecker.ValidTrainingImage>,
        onProgress: (Int, Int) -> Unit = { _, _ -> }
    ): String = withContext(Dispatchers.Default) {

        val person = personDao.getPersonById(personId)
            ?: throw IllegalArgumentException("Person with ID $personId not found")

        val embeddings = faceNetModel.generateEmbeddingsBatch(
            faceBitmaps = validImages.map { it.croppedFaceBitmap },
            onProgress = onProgress
        )

        val personEmbedding = faceNetModel.createPersonModel(embeddings)

        val confidences = embeddings.map { embedding ->
            faceNetModel.calculateSimilarity(personEmbedding, embedding)
        }
        val avgConfidence = confidences.average().toFloat()

        // Compute distribution stats for self-calibrating rejection
        val stdDev = kotlin.math.sqrt(
            confidences.map { (it - avgConfidence).toDouble().let { d -> d * d } }.average()
        ).toFloat()
        val minSimilarity = confidences.minOrNull() ?: 0f

        val faceModel = FaceModelEntity.create(
            personId = personId,
            embeddingArray = personEmbedding,
            trainingImageCount = validImages.size,
            averageConfidence = avgConfidence,
            similarityStdDev = stdDev,
            similarityMin = minSimilarity
        )

        faceModelDao.insertFaceModel(faceModel)
        faceModel.id
    }

    /**
     * Retrain face model with additional images.
     */
    suspend fun retrainFaceModel(
        faceModelId: String,
        newFaceImages: List<Bitmap>
    ) = withContext(Dispatchers.Default) {

        val faceModel = faceModelDao.getFaceModelById(faceModelId)
            ?: throw IllegalArgumentException("Face model $faceModelId not found")

        val existingEmbedding = faceModel.getEmbeddingArray()
        val newEmbeddings = faceNetModel.generateEmbeddingsBatch(newFaceImages)
        val allEmbeddings = listOf(existingEmbedding) + newEmbeddings
        val updatedEmbedding = faceNetModel.createPersonModel(allEmbeddings)

        val confidences = allEmbeddings.map { embedding ->
            faceNetModel.calculateSimilarity(updatedEmbedding, embedding)
        }
        val avgConfidence = confidences.average().toFloat()

        faceModelDao.updateFaceModel(
            FaceModelEntity.create(
                personId = faceModel.personId,
                embeddingArray = updatedEmbedding,
                trainingImageCount = faceModel.trainingImageCount + newFaceImages.size,
                averageConfidence = avgConfidence
            ).copy(
                id = faceModelId,
                createdAt = faceModel.createdAt,
                updatedAt = System.currentTimeMillis()
            )
        )
    }

    // ======================
    // SCANNING / RECOGNITION
    // ======================

    /**
     * Scan an image for faces and tag recognized persons.
     *
     * ALSO UPDATES FACE DETECTION CACHE for optimization.
     *
     * @param imageId String (from ImageEntity.imageId)
     */
    suspend fun scanImage(
        imageId: String,
        detectedFaces: List<DetectedFace>,
        threshold: Float = FaceNetModel.SIMILARITY_THRESHOLD_HIGH
    ): List<PhotoFaceTagEntity> = withContext(Dispatchers.Default) {

        // OPTIMIZATION: Update face detection cache
        // This makes future scans faster by skipping images without faces
        withContext(Dispatchers.IO) {
            imageDao.updateFaceDetectionCache(
                imageId = imageId,
                hasFaces = detectedFaces.isNotEmpty(),
                faceCount = detectedFaces.size
            )
        }

        val faceModels = faceModelDao.getAllActiveFaceModels()

        if (faceModels.isEmpty()) {
            return@withContext emptyList()
        }

        val tags = mutableListOf<PhotoFaceTagEntity>()

        for (detectedFace in detectedFaces) {
            val faceEmbedding = faceNetModel.generateEmbedding(detectedFace.croppedBitmap)

            var bestMatch: Pair<String, Float>? = null
            var highestSimilarity = threshold

            for (faceModel in faceModels) {
                // Check ALL centroids for best match (critical for children with age centroids)
                val centroids = faceModel.getCentroids()
                val bestCentroidSimilarity = centroids.maxOfOrNull { centroid ->
                    faceNetModel.calculateSimilarity(faceEmbedding, centroid.getEmbeddingArray())
                } ?: 0f

                if (bestCentroidSimilarity > highestSimilarity) {
                    highestSimilarity = bestCentroidSimilarity
                    bestMatch = Pair(faceModel.id, bestCentroidSimilarity)
                }
            }

            if (bestMatch != null) {
                val (faceModelId, confidence) = bestMatch

                val tag = PhotoFaceTagEntity.create(
                    imageId = imageId,
                    faceModelId = faceModelId,
                    boundingBox = detectedFace.boundingBox,
                    confidence = confidence,
                    faceEmbedding = faceEmbedding
                )

                tags.add(tag)
                faceModelDao.updateLastUsed(faceModelId, System.currentTimeMillis())
            }
        }

        if (tags.isNotEmpty()) {
            photoFaceTagDao.insertTags(tags)
        }

        tags
    }

    /**
     * Recognize a single face bitmap (without saving).
     */
    suspend fun recognizeFace(
        faceBitmap: Bitmap,
        threshold: Float = FaceNetModel.SIMILARITY_THRESHOLD_HIGH
    ): Pair<String, Float>? = withContext(Dispatchers.Default) {

        val faceEmbedding = faceNetModel.generateEmbedding(faceBitmap)
        val faceModels = faceModelDao.getAllActiveFaceModels()
        val modelEmbeddings = faceModels.map { it.id to it.getEmbeddingArray() }

        faceNetModel.findBestMatch(faceEmbedding, modelEmbeddings, threshold)
    }

    // ======================
    // SEARCH / QUERY
    // ======================

    /**
     * Get all images containing a specific person.
     *
     * @param personId String UUID
     */
    suspend fun getImagesForPerson(personId: String): List<ImageEntity> = withContext(Dispatchers.IO) {

        val faceModel = faceModelDao.getFaceModelByPersonId(personId)
            ?: return@withContext emptyList()

        val imageIds = photoFaceTagDao.getImageIdsForFaceModel(faceModel.id)
        imageDao.getImagesByIds(imageIds)
    }

    /**
     * Get images for person as Flow (reactive).
     */
    fun getImagesForPersonFlow(personId: String): Flow<List<ImageEntity>> {
        return photoFaceTagDao.getImageIdsForFaceModelFlow(personId)
            .map { imageIds ->
                imageDao.getImagesByIds(imageIds)
            }
    }

    /**
     * Get all persons with face models.
     */
    suspend fun getPersonsWithFaceModels(): List<PersonEntity> = withContext(Dispatchers.IO) {
        val faceModels = faceModelDao.getAllActiveFaceModels()
        val personIds = faceModels.map { it.personId }
        personDao.getPersonsByIds(personIds)
    }

    /**
     * Get face detection stats for a person.
     */
    suspend fun getPersonFaceStats(personId: String): PersonFaceStats? = withContext(Dispatchers.IO) {

        val person = personDao.getPersonById(personId) ?: return@withContext null
        val faceModel = faceModelDao.getFaceModelByPersonId(personId) ?: return@withContext null

        val imageIds = photoFaceTagDao.getImageIdsForFaceModel(faceModel.id)
        val allTags = photoFaceTagDao.getAllTagsForFaceModel(faceModel.id)

        val avgConfidence = if (allTags.isNotEmpty()) {
            allTags.map { it.confidence }.average().toFloat()
        } else {
            0f
        }
        val lastDetected = allTags.maxOfOrNull { it.detectedAt }

        PersonFaceStats(
            personId = person.id,
            personName = person.name,
            faceModelId = faceModel.id,
            trainingImageCount = faceModel.trainingImageCount,
            taggedPhotoCount = imageIds.size,
            averageConfidence = avgConfidence,
            lastDetectedAt = lastDetected
        )
    }

    /**
     * Get face tags for an image.
     */
    suspend fun getFaceTagsForImage(imageId: String): List<PhotoFaceTagEntity> {
        return photoFaceTagDao.getTagsForImage(imageId)
    }

    /**
     * Get person from a face tag.
     */
    suspend fun getPersonForFaceTag(tag: PhotoFaceTagEntity): PersonEntity? = withContext(Dispatchers.IO) {
        val faceModel = faceModelDao.getFaceModelById(tag.faceModelId) ?: return@withContext null
        personDao.getPersonById(faceModel.personId)
    }

    /**
     * Get face tags with person info for an image.
     */
    suspend fun getFaceTagsWithPersons(imageId: String): List<Pair<PhotoFaceTagEntity, PersonEntity>> = withContext(Dispatchers.IO) {
        val tags = photoFaceTagDao.getTagsForImage(imageId)
        tags.mapNotNull { tag ->
            val person = getPersonForFaceTag(tag)
            if (person != null) tag to person else null
        }
    }

    // ======================
    // VERIFICATION / QUALITY
    // ======================

    suspend fun verifyFaceTag(tagId: String) {
        photoFaceTagDao.markTagAsVerified(
            tagId = tagId,
            timestamp = System.currentTimeMillis()
        )
    }

    suspend fun getUnverifiedTags(): List<PhotoFaceTagEntity> {
        return photoFaceTagDao.getUnverifiedTags()
    }

    suspend fun getLowConfidenceTags(threshold: Float = 0.7f): List<PhotoFaceTagEntity> {
        return photoFaceTagDao.getLowConfidenceTags(threshold)
    }

    // ======================
    // MANAGEMENT
    // ======================

    suspend fun deleteFaceModel(faceModelId: String) = withContext(Dispatchers.IO) {
        photoFaceTagDao.deleteTagsForFaceModel(faceModelId)
        faceModelDao.deleteFaceModelById(faceModelId)
    }

    // Add this method to FaceRecognitionRepository_StringIds.kt
// Replace the existing createPersonWithFaceModel method with this version:

    /**
     * Create a new person with face model in one operation.
     * Now supports full PersonEntity with optional fields.
     *
     * @param person PersonEntity with name, DOB, relationship, etc.
     * @return PersonId (String UUID)
     */
    suspend fun createPersonWithFaceModel(
        person: PersonEntity,
        validImages: List<TrainingSanityChecker.ValidTrainingImage>,
        onProgress: (Int, Int) -> Unit = { _, _ -> }
    ): String = withContext(Dispatchers.IO) {

        // Insert person with all fields
        personDao.insert(person)

        // Train face model
        trainPerson(
            personId = person.id,
            validImages = validImages,
            onProgress = onProgress
        )

        // Generate age tags for children
        if (person.isChild && person.dateOfBirth != null) {
            generateAgeTagsForTraining(person, validImages)
        }

        person.id
    }

    /**
     * Generate age tags from training images for a child
     */
    private suspend fun generateAgeTagsForTraining(
        person: PersonEntity,
        validImages: List<TrainingSanityChecker.ValidTrainingImage>
    ) {
        try {
            val dob = person.dateOfBirth ?: return

            val tags = validImages.mapNotNull { img ->
                val imageEntity = imageDao.getImageByUri(img.uri.toString()) ?: return@mapNotNull null
                val ageMs = imageEntity.capturedAt - dob
                val ageYears = (ageMs / (365.25 * 24 * 60 * 60 * 1000)).toInt()

                if (ageYears < 0 || ageYears > 25) return@mapNotNull null

                PersonAgeTagEntity.create(
                    personId = person.id,
                    personName = person.name,
                    imageId = imageEntity.imageId,
                    ageAtCapture = ageYears,
                    confidence = 1.0f
                )
            }

            if (tags.isNotEmpty()) {
                personAgeTagDao.insertTags(tags)
                Log.d(TAG, "Created ${tags.size} age tags for ${person.name}")
            }
        } catch (e: Exception) {
            Log.e(TAG, "Failed to generate age tags", e)
        }
    }

    /**
     * Get face model by ID
     */
    suspend fun getFaceModelById(faceModelId: String): FaceModelEntity? = withContext(Dispatchers.IO) {
        faceModelDao.getFaceModelById(faceModelId)
    }

    suspend fun deleteTagsForImage(imageId: String) {
        photoFaceTagDao.deleteTagsForImage(imageId)
    }

    /**
     * Get all image IDs that have been tagged with this face model
     * Used for scan optimization (skip already-tagged images)
     */
    suspend fun getImageIdsForFaceModel(faceModelId: String): List<String> = withContext(Dispatchers.IO) {
        photoFaceTagDao.getImageIdsForFaceModel(faceModelId)
    }

    fun cleanup() {
        faceNetModel.close()
    }


}

data class DetectedFace(
    val croppedBitmap: Bitmap,
    val boundingBox: android.graphics.Rect
)

data class PersonFaceStats(
    val personId: String,
    val personName: String,
    val faceModelId: String,
    val trainingImageCount: Int,
    val taggedPhotoCount: Int,
    val averageConfidence: Float,
    val lastDetectedAt: Long?
)