package com.placeholder.sherpai2.ml

import android.content.Context
import android.graphics.Bitmap
import android.util.Log
import org.tensorflow.lite.Interpreter
import java.io.FileInputStream
import java.nio.ByteBuffer
import java.nio.ByteOrder
import java.nio.MappedByteBuffer
import java.nio.channels.FileChannel
import kotlin.math.sqrt

/**
 * FaceNetModel - MobileFaceNet wrapper with debugging
 *
 * IMPROVEMENTS:
 * - ✅ Detailed error logging
 * - ✅ Model validation on init
 * - ✅ Embedding validation (detect all-zeros)
 * - ✅ Toggle-able debug mode
 */
class FaceNetModel(
    private val context: Context,
    private val debugMode: Boolean = true  // Enable for troubleshooting
) {

    companion object {
        private const val TAG = "FaceNetModel"
        private const val MODEL_FILE = "mobilefacenet.tflite"
        private const val INPUT_SIZE = 112
        private const val EMBEDDING_SIZE = 192

        const val SIMILARITY_THRESHOLD_HIGH = 0.7f
        const val SIMILARITY_THRESHOLD_MEDIUM = 0.6f
        const val SIMILARITY_THRESHOLD_LOW = 0.5f
    }

    private var interpreter: Interpreter? = null
    private var modelLoadSuccess = false

    init {
        try {
            if (debugMode) Log.d(TAG, "Loading FaceNet model: $MODEL_FILE")

            val model = loadModelFile()
            interpreter = Interpreter(model)
            modelLoadSuccess = true

            if (debugMode) {
                Log.d(TAG, "✅ FaceNet model loaded successfully")
                Log.d(TAG, "Model input size: ${INPUT_SIZE}x$INPUT_SIZE")
                Log.d(TAG, "Embedding size: $EMBEDDING_SIZE")
            }

            // Test model with dummy input
            testModel()

        } catch (e: Exception) {
            Log.e(TAG, "❌ CRITICAL: Failed to load FaceNet model from assets/$MODEL_FILE", e)
            Log.e(TAG, "Make sure mobilefacenet.tflite exists in app/src/main/assets/")
            modelLoadSuccess = false
            throw RuntimeException("Failed to load FaceNet model: ${e.message}", e)
        }
    }

    /**
     * Test model with dummy input to verify it works
     */
    private fun testModel() {
        try {
            val testBitmap = Bitmap.createBitmap(INPUT_SIZE, INPUT_SIZE, Bitmap.Config.ARGB_8888)
            val testEmbedding = generateEmbedding(testBitmap)
            testBitmap.recycle()

            val sum = testEmbedding.sum()
            val norm = sqrt(testEmbedding.map { it * it }.sum())

            if (debugMode) {
                Log.d(TAG, "Model test: embedding sum=$sum, norm=$norm")
            }

            if (sum == 0f || norm == 0f) {
                Log.e(TAG, "⚠️ WARNING: Model test produced zero embedding!")
            } else {
                if (debugMode) Log.d(TAG, "✅ Model test passed")
            }
        } catch (e: Exception) {
            Log.e(TAG, "Model test failed", e)
        }
    }

    /**
     * Load TFLite model from assets
     */
    private fun loadModelFile(): MappedByteBuffer {
        try {
            val fileDescriptor = context.assets.openFd(MODEL_FILE)
            val inputStream = FileInputStream(fileDescriptor.fileDescriptor)
            val fileChannel = inputStream.channel
            val startOffset = fileDescriptor.startOffset
            val declaredLength = fileDescriptor.declaredLength

            if (debugMode) {
                Log.d(TAG, "Model file size: ${declaredLength / 1024}KB")
            }

            return fileChannel.map(FileChannel.MapMode.READ_ONLY, startOffset, declaredLength)
        } catch (e: Exception) {
            Log.e(TAG, "Failed to open model file: $MODEL_FILE", e)
            throw e
        }
    }

    /**
     * Generate embedding for a single face
     *
     * @param faceBitmap Cropped face image (will be resized to 112x112)
     * @return 192-dimensional embedding
     */
    fun generateEmbedding(faceBitmap: Bitmap): FloatArray {
        if (!modelLoadSuccess || interpreter == null) {
            Log.e(TAG, "❌ Cannot generate embedding: model not loaded!")
            return FloatArray(EMBEDDING_SIZE) { 0f }
        }

        try {
            val resized = Bitmap.createScaledBitmap(faceBitmap, INPUT_SIZE, INPUT_SIZE, true)
            val inputBuffer = preprocessImage(resized)
            val output = Array(1) { FloatArray(EMBEDDING_SIZE) }

            interpreter?.run(inputBuffer, output)

            val normalized = normalizeEmbedding(output[0])

            // DIAGNOSTIC: Check embedding quality
            if (debugMode) {
                val sum = normalized.sum()
                val norm = sqrt(normalized.map { it * it }.sum())

                if (sum == 0f && norm == 0f) {
                    Log.e(TAG, "❌ CRITICAL: Generated all-zero embedding!")
                    Log.e(TAG, "Input bitmap: ${faceBitmap.width}x${faceBitmap.height}")
                } else {
                    Log.d(TAG, "✅ Embedding: sum=${"%.2f".format(sum)}, norm=${"%.2f".format(norm)}, first5=[${normalized.take(5).joinToString { "%.3f".format(it) }}]")
                }
            }

            return normalized

        } catch (e: Exception) {
            Log.e(TAG, "Failed to generate embedding", e)
            return FloatArray(EMBEDDING_SIZE) { 0f }
        }
    }

    /**
     * Generate embeddings for multiple faces (batch processing)
     */
    fun generateEmbeddingsBatch(
        faceBitmaps: List<Bitmap>,
        onProgress: (Int, Int) -> Unit = { _, _ -> }
    ): List<FloatArray> {
        if (debugMode) {
            Log.d(TAG, "Generating embeddings for ${faceBitmaps.size} faces")
        }

        return faceBitmaps.mapIndexed { index, bitmap ->
            onProgress(index + 1, faceBitmaps.size)
            generateEmbedding(bitmap)
        }
    }

    /**
     * Create person model by averaging multiple embeddings
     */
    fun createPersonModel(embeddings: List<FloatArray>): FloatArray {
        require(embeddings.isNotEmpty()) { "Need at least one embedding" }

        if (debugMode) {
            Log.d(TAG, "Creating person model from ${embeddings.size} embeddings")
        }

        val averaged = FloatArray(EMBEDDING_SIZE) { 0f }

        embeddings.forEach { embedding ->
            for (i in embedding.indices) {
                averaged[i] += embedding[i]
            }
        }

        val count = embeddings.size.toFloat()
        for (i in averaged.indices) {
            averaged[i] /= count
        }

        val normalized = normalizeEmbedding(averaged)

        if (debugMode) {
            val sum = normalized.sum()
            Log.d(TAG, "Person model created: sum=${"%.2f".format(sum)}")
        }

        return normalized
    }

    /**
     * Calculate cosine similarity between two embeddings
     * Returns value between -1.0 and 1.0 (higher = more similar)
     */
    fun calculateSimilarity(embedding1: FloatArray, embedding2: FloatArray): Float {
        require(embedding1.size == EMBEDDING_SIZE && embedding2.size == EMBEDDING_SIZE) {
            "Invalid embedding size: ${embedding1.size} vs ${embedding2.size}"
        }

        var dotProduct = 0f
        var norm1 = 0f
        var norm2 = 0f

        for (i in embedding1.indices) {
            dotProduct += embedding1[i] * embedding2[i]
            norm1 += embedding1[i] * embedding1[i]
            norm2 += embedding2[i] * embedding2[i]
        }

        val similarity = dotProduct / (sqrt(norm1) * sqrt(norm2))

        if (debugMode && (similarity.isNaN() || similarity.isInfinite())) {
            Log.e(TAG, "❌ Invalid similarity: $similarity (norm1=$norm1, norm2=$norm2)")
            return 0f
        }

        return similarity
    }

    /**
     * Find best matching face model from a list
     *
     * @param faceEmbedding Embedding to match
     * @param modelEmbeddings List of (modelId: String, embedding: FloatArray)
     * @param threshold Minimum similarity threshold
     * @return Pair of (modelId: String, confidence: Float) or null
     */
    fun findBestMatch(
        faceEmbedding: FloatArray,
        modelEmbeddings: List<Pair<String, FloatArray>>,
        threshold: Float = SIMILARITY_THRESHOLD_HIGH
    ): Pair<String, Float>? {
        var bestMatch: Pair<String, Float>? = null
        var highestSimilarity = threshold

        for ((modelId, modelEmbedding) in modelEmbeddings) {
            val similarity = calculateSimilarity(faceEmbedding, modelEmbedding)

            if (similarity > highestSimilarity) {
                highestSimilarity = similarity
                bestMatch = Pair(modelId, similarity)
            }
        }

        if (debugMode && bestMatch != null) {
            Log.d(TAG, "Best match: ${bestMatch.first} with similarity ${bestMatch.second}")
        }

        return bestMatch
    }

    /**
     * Preprocess image for model input
     */
    private fun preprocessImage(bitmap: Bitmap): ByteBuffer {
        val buffer = ByteBuffer.allocateDirect(4 * INPUT_SIZE * INPUT_SIZE * 3)
        buffer.order(ByteOrder.nativeOrder())

        val pixels = IntArray(INPUT_SIZE * INPUT_SIZE)
        bitmap.getPixels(pixels, 0, INPUT_SIZE, 0, 0, INPUT_SIZE, INPUT_SIZE)

        for (pixel in pixels) {
            val r = ((pixel shr 16) and 0xFF) / 255.0f
            val g = ((pixel shr 8) and 0xFF) / 255.0f
            val b = (pixel and 0xFF) / 255.0f

            // Normalize to [-1, 1]
            buffer.putFloat((r - 0.5f) / 0.5f)
            buffer.putFloat((g - 0.5f) / 0.5f)
            buffer.putFloat((b - 0.5f) / 0.5f)
        }

        return buffer
    }

    /**
     * Normalize embedding to unit length
     */
    private fun normalizeEmbedding(embedding: FloatArray): FloatArray {
        var norm = 0f
        for (value in embedding) {
            norm += value * value
        }
        norm = sqrt(norm)

        return if (norm > 0) {
            FloatArray(embedding.size) { i -> embedding[i] / norm }
        } else {
            Log.w(TAG, "⚠️ Cannot normalize zero embedding")
            embedding
        }
    }

    /**
     * Get model status for diagnostics
     */
    fun getModelStatus(): String {
        return if (modelLoadSuccess) {
            "✅ Model loaded and operational"
        } else {
            "❌ Model failed to load - check assets/$MODEL_FILE"
        }
    }

    /**
     * Clean up resources
     */
    fun close() {
        if (debugMode) {
            Log.d(TAG, "Closing FaceNet model")
        }
        interpreter?.close()
        interpreter = null
    }
}