package com.placeholder.sherpai2.ui.trainingprep

import android.content.Context
import android.graphics.Bitmap
import android.net.Uri

/**
 * ENHANCED TrainingSanityChecker
 *
 * New features:
 * - Progress callbacks
 * - Exclude functionality
 * - Faster processing
 */
class TrainingSanityChecker(private val context: Context) {

    private val faceDetectionHelper = FaceDetectionHelper(context)
    private val duplicateDetector = DuplicateImageDetector(context)

    data class SanityCheckResult(
        val isValid: Boolean,
        val faceDetectionResults: List<FaceDetectionHelper.FaceDetectionResult>,
        val duplicateCheckResult: DuplicateImageDetector.DuplicateCheckResult,
        val validationErrors: List<ValidationError>,
        val warnings: List<String>,
        val validImagesWithFaces: List<ValidTrainingImage>,
        val excludedImages: Set<Uri> = emptySet() // NEW: Track excluded images
    )

    data class ValidTrainingImage(
        val uri: Uri,
        val croppedFaceBitmap: Bitmap,
        val faceCount: Int
    )

    sealed class ValidationError {
        data class NoFaceDetected(val uris: List<Uri>) : ValidationError()
        data class MultipleFacesDetected(val uri: Uri, val faceCount: Int) : ValidationError()
        data class DuplicateImages(val groups: List<DuplicateImageDetector.DuplicateGroup>) : ValidationError()
        data class InsufficientImages(val required: Int, val available: Int) : ValidationError()
        data class ImageLoadError(val uri: Uri, val error: String) : ValidationError()
    }

    /**
     * Perform comprehensive sanity checks with PROGRESS tracking
     */
    suspend fun performSanityChecks(
        imageUris: List<Uri>,
        minImagesRequired: Int = 15,
        allowMultipleFaces: Boolean = false,
        duplicateSimilarityThreshold: Double = 0.95,
        excludedImages: Set<Uri> = emptySet(), // NEW: Allow excluding images
        onProgress: ((String, Int, Int) -> Unit)? = null // NEW: Progress callback
    ): SanityCheckResult {

        val validationErrors = mutableListOf<ValidationError>()
        val warnings = mutableListOf<String>()

        // Filter out excluded images
        val activeImages = imageUris.filter { it !in excludedImages }

        // Check minimum image count (AFTER exclusions)
        if (activeImages.size < minImagesRequired) {
            validationErrors.add(
                ValidationError.InsufficientImages(
                    required = minImagesRequired,
                    available = activeImages.size
                )
            )
        }

        // Step 1: Detect faces in all images (WITH PROGRESS)
        onProgress?.invoke("Detecting faces...", 0, activeImages.size)

        val faceDetectionResults = faceDetectionHelper.detectFacesInImages(
            uris = activeImages,
            onProgress = { current, total ->
                onProgress?.invoke("Detecting faces...", current, total)
            }
        )

        // Check for images without faces
        val imagesWithoutFaces = faceDetectionResults.filter { !it.hasFace }
        if (imagesWithoutFaces.isNotEmpty()) {
            validationErrors.add(
                ValidationError.NoFaceDetected(
                    uris = imagesWithoutFaces.map { it.uri }
                )
            )
        }

        // Check for images with errors
        faceDetectionResults.filter { it.errorMessage != null }.forEach { result ->
            validationErrors.add(
                ValidationError.ImageLoadError(
                    uri = result.uri,
                    error = result.errorMessage ?: "Unknown error"
                )
            )
        }

        // Check for images with multiple faces
        if (!allowMultipleFaces) {
            faceDetectionResults.filter { it.faceCount > 1 }.forEach { result ->
                validationErrors.add(
                    ValidationError.MultipleFacesDetected(
                        uri = result.uri,
                        faceCount = result.faceCount
                    )
                )
            }
        } else {
            faceDetectionResults.filter { it.faceCount > 1 }.forEach { result ->
                warnings.add("Image ${result.uri.lastPathSegment} contains ${result.faceCount} faces. Using the largest detected face.")
            }
        }

        // Step 2: Check for duplicate images
        onProgress?.invoke("Checking for duplicates...", activeImages.size, activeImages.size)

        val duplicateCheckResult = duplicateDetector.checkForDuplicates(
            uris = activeImages,
            similarityThreshold = duplicateSimilarityThreshold
        )

        if (duplicateCheckResult.hasDuplicates) {
            validationErrors.add(
                ValidationError.DuplicateImages(
                    groups = duplicateCheckResult.duplicateGroups
                )
            )
        }

        // Step 3: Create list of valid training images
        val validImagesWithFaces = faceDetectionResults
            .filter { it.hasFace && it.croppedFaceBitmap != null }
            .filter { allowMultipleFaces || it.faceCount == 1 }
            .map { result ->
                ValidTrainingImage(
                    uri = result.uri,
                    croppedFaceBitmap = result.croppedFaceBitmap!!,
                    faceCount = result.faceCount
                )
            }

        // Check if we have enough valid images after all checks
        if (validImagesWithFaces.size < minImagesRequired) {
            val existingError = validationErrors.find { it is ValidationError.InsufficientImages }
            if (existingError == null) {
                validationErrors.add(
                    ValidationError.InsufficientImages(
                        required = minImagesRequired,
                        available = validImagesWithFaces.size
                    )
                )
            }
        }

        val isValid = validationErrors.isEmpty() && validImagesWithFaces.size >= minImagesRequired

        onProgress?.invoke("Analysis complete", activeImages.size, activeImages.size)

        return SanityCheckResult(
            isValid = isValid,
            faceDetectionResults = faceDetectionResults,
            duplicateCheckResult = duplicateCheckResult,
            validationErrors = validationErrors,
            warnings = warnings,
            validImagesWithFaces = validImagesWithFaces,
            excludedImages = excludedImages
        )
    }

    /**
     * Format validation errors into human-readable messages
     */
    fun formatValidationErrors(errors: List<ValidationError>): List<String> {
        return errors.map { error ->
            when (error) {
                is ValidationError.NoFaceDetected -> {
                    val count = error.uris.size
                    "No face detected in $count image(s)"
                }
                is ValidationError.MultipleFacesDetected -> {
                    "Multiple faces (${error.faceCount}) detected in: ${error.uri.lastPathSegment}"
                }
                is ValidationError.DuplicateImages -> {
                    val count = error.groups.size
                    "Found $count duplicate group(s)"
                }
                is ValidationError.InsufficientImages -> {
                    "Need ${error.required} images, have ${error.available}"
                }
                is ValidationError.ImageLoadError -> {
                    "Failed to load image: ${error.uri.lastPathSegment}"
                }
            }
        }
    }

    /**
     * Clean up resources
     */
    fun cleanup() {
        faceDetectionHelper.cleanup()
    }
}