package com.placeholder.sherpai2.ui.trainingprep

import android.app.Application
import android.graphics.Bitmap
import android.net.Uri
import androidx.lifecycle.AndroidViewModel
import androidx.lifecycle.viewModelScope
import androidx.datastore.preferences.core.booleanPreferencesKey
import androidx.datastore.preferences.preferencesDataStore
import androidx.work.WorkManager
import android.content.Context
import com.placeholder.sherpai2.data.local.entity.PersonEntity
import com.placeholder.sherpai2.data.repository.FaceRecognitionRepository
import com.placeholder.sherpai2.ml.FaceNetModel
import com.placeholder.sherpai2.workers.LibraryScanWorker
import dagger.hilt.android.lifecycle.HiltViewModel
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.first
import kotlinx.coroutines.flow.map
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.launch
import javax.inject.Inject

sealed class ScanningState {
    object Idle : ScanningState()
    data class Processing(val progress: Int, val total: Int) : ScanningState()
    data class Success(
        val sanityCheckResult: TrainingSanityChecker.SanityCheckResult
    ) : ScanningState()
    data class Error(val message: String) : ScanningState()
}

sealed class TrainingState {
    object Idle : TrainingState()
    data class Processing(val stage: String, val progress: Int, val total: Int) : TrainingState()
    data class Success(
        val personName: String,
        val personId: String,
        val relationship: String?
    ) : TrainingState()
    data class Error(val message: String) : TrainingState()
}

/**
 * Person info captured before photo selection
 */
data class PersonInfo(
    val name: String,
    val dateOfBirth: Long?,
    val relationship: String,
    val isChild: Boolean = false
)

/**
 * FIXED TrainViewModel with proper exclude functionality and efficient replace
 */
private val android.content.Context.dataStore by preferencesDataStore(name = "settings")
private val KEY_BACKGROUND_TAGGING = booleanPreferencesKey("background_recognition_tagging")

@HiltViewModel
class TrainViewModel @Inject constructor(
    application: Application,
    private val faceRecognitionRepository: FaceRecognitionRepository,
    private val faceNetModel: FaceNetModel,
    private val workManager: WorkManager
) : AndroidViewModel(application) {

    private val sanityChecker = TrainingSanityChecker(application)
    private val faceDetectionHelper = FaceDetectionHelper(application)
    private val dataStore = application.dataStore

    private val _uiState = MutableStateFlow<ScanningState>(ScanningState.Idle)
    val uiState: StateFlow<ScanningState> = _uiState.asStateFlow()

    private val _trainingState = MutableStateFlow<TrainingState>(TrainingState.Idle)
    val trainingState: StateFlow<TrainingState> = _trainingState.asStateFlow()

    // Store person info for later use during training
    private var personInfo: PersonInfo? = null

    private var currentImageUris: List<Uri> = emptyList()
    private val manualFaceSelections = mutableMapOf<Uri, ManualFaceSelection>()

    // Track excluded images
    private val excludedImages = mutableSetOf<Uri>()

    data class ManualFaceSelection(
        val faceIndex: Int,
        val croppedFaceBitmap: Bitmap
    )

    /**
     * Store person info before photo selection
     */
    fun setPersonInfo(name: String, dateOfBirth: Long?, relationship: String, isChild: Boolean = false) {
        personInfo = PersonInfo(name, dateOfBirth, relationship, isChild)
    }

    /**
     * Get stored person info
     */
    fun getPersonInfo(): PersonInfo? = personInfo

    /**
     * Exclude an image from training
     */
    fun excludeImage(uri: Uri) {
        excludedImages.add(uri)

        val currentState = _uiState.value
        if (currentState is ScanningState.Success) {
            val updatedResult = applyManualSelections(currentState.sanityCheckResult)
            _uiState.value = ScanningState.Success(updatedResult)
        }
    }

    /**
     * Include a previously excluded image
     */
    fun includeImage(uri: Uri) {
        excludedImages.remove(uri)

        val currentState = _uiState.value
        if (currentState is ScanningState.Success) {
            val updatedResult = applyManualSelections(currentState.sanityCheckResult)
            _uiState.value = ScanningState.Success(updatedResult)
        }
    }

    /**
     * Check if an image is excluded
     */
    fun isImageExcluded(uri: Uri): Boolean {
        return uri in excludedImages
    }

    /**
     * Create face model with captured person info
     */
    fun createFaceModel(personName: String) {
        val currentState = _uiState.value
        if (currentState !is ScanningState.Success) {
            _trainingState.value = TrainingState.Error("No validated images available")
            return
        }

        val validImages = currentState.sanityCheckResult.validImagesWithFaces
        if (validImages.size < 15) {
            _trainingState.value = TrainingState.Error(
                "Need at least 15 valid images, have ${validImages.size}"
            )
            return
        }

        viewModelScope.launch {
            try {
                _trainingState.value = TrainingState.Processing(
                    stage = "Creating person and training model",
                    progress = 0,
                    total = validImages.size
                )

                val person = PersonEntity.create(
                    name = personName,
                    dateOfBirth = personInfo?.dateOfBirth,
                    isChild = personInfo?.isChild ?: false,
                    relationship = personInfo?.relationship
                )

                val personId = faceRecognitionRepository.createPersonWithFaceModel(
                    person = person,
                    validImages = validImages,
                    onProgress = { current, total ->
                        _trainingState.value = TrainingState.Processing(
                            stage = "Processing image $current/$total",
                            progress = current,
                            total = total
                        )
                    }
                )

                _trainingState.value = TrainingState.Success(
                    personName = personName,
                    personId = personId,
                    relationship = person.relationship
                )

                // Trigger library scan if setting enabled
                val backgroundTaggingEnabled = dataStore.data
                    .map { it[KEY_BACKGROUND_TAGGING] ?: true }
                    .first()

                if (backgroundTaggingEnabled) {
                    // Use default threshold (0.62 solo, 0.68 group)
                    val scanRequest = LibraryScanWorker.createWorkRequest(
                        personId = personId,
                        personName = personName
                    )
                    workManager.enqueue(scanRequest)
                }

            } catch (e: Exception) {
                _trainingState.value = TrainingState.Error(
                    e.message ?: "Failed to create face model"
                )
            }
        }
    }

    fun resetTrainingState() {
        _trainingState.value = TrainingState.Idle
    }

    fun scanAndTagFaces(imageUris: List<Uri>) {
        currentImageUris = imageUris
        manualFaceSelections.clear()
        excludedImages.clear()
        performScan(imageUris)
    }

    /**
     * FIXED: Replace image - only rescan the ONE new image, not all images!
     */
    fun replaceImage(oldUri: Uri, newUri: Uri) {
        viewModelScope.launch {
            try {
                val currentState = _uiState.value
                if (currentState !is ScanningState.Success) return@launch

                // Update the URI list
                val updatedUris = currentImageUris.toMutableList()
                val index = updatedUris.indexOf(oldUri)
                if (index == -1) return@launch

                updatedUris[index] = newUri
                currentImageUris = updatedUris

                // Clean up old selections/exclusions
                manualFaceSelections.remove(oldUri)
                excludedImages.remove(oldUri)

                // Only scan the NEW image
                val newResult = faceDetectionHelper.detectFacesInImage(newUri)

                // Update the results list
                val updatedFaceResults = currentState.sanityCheckResult.faceDetectionResults.toMutableList()
                updatedFaceResults[index] = newResult

                // Create updated SanityCheckResult
                val updatedSanityResult = currentState.sanityCheckResult.copy(
                    faceDetectionResults = updatedFaceResults
                )

                // Apply manual selections and exclusions
                val finalResult = applyManualSelections(updatedSanityResult)
                _uiState.value = ScanningState.Success(finalResult)

            } catch (e: Exception) {
                _uiState.value = ScanningState.Error(
                    e.message ?: "Failed to replace image"
                )
            }
        }
    }

    /**
     * Select face and auto-include the image
     */
    fun selectFaceFromImage(imageUri: Uri, faceIndex: Int, croppedFaceBitmap: Bitmap) {
        manualFaceSelections[imageUri] = ManualFaceSelection(faceIndex, croppedFaceBitmap)
        excludedImages.remove(imageUri) // Auto-include

        val currentState = _uiState.value
        if (currentState is ScanningState.Success) {
            val updatedResult = applyManualSelections(currentState.sanityCheckResult)
            _uiState.value = ScanningState.Success(updatedResult)
        }
    }

    /**
     * Perform full scan with exclusions and progress tracking
     */
    private fun performScan(imageUris: List<Uri>) {
        viewModelScope.launch {
            try {
                _uiState.value = ScanningState.Processing(0, imageUris.size)

                val result = sanityChecker.performSanityChecks(
                    imageUris = imageUris,
                    minImagesRequired = 15,
                    allowMultipleFaces = true,
                    duplicateSimilarityThreshold = 0.95,
                    excludedImages = excludedImages,
                    onProgress = { stage, current, total ->
                        _uiState.value = ScanningState.Processing(current, total)
                    }
                )

                val finalResult = applyManualSelections(result)
                _uiState.value = ScanningState.Success(finalResult)

            } catch (e: Exception) {
                _uiState.value = ScanningState.Error(
                    e.message ?: "An unknown error occurred"
                )
            }
        }
    }

    /**
     * Apply manual selections with exclusion filtering
     */
    private fun applyManualSelections(
        result: TrainingSanityChecker.SanityCheckResult
    ): TrainingSanityChecker.SanityCheckResult {

        if (manualFaceSelections.isEmpty() && excludedImages.isEmpty()) {
            return result
        }

        val updatedFaceResults = result.faceDetectionResults.map { faceResult ->
            val manualSelection = manualFaceSelections[faceResult.uri]
            if (manualSelection != null) {
                faceResult.copy(
                    croppedFaceBitmap = manualSelection.croppedFaceBitmap,
                    faceCount = 1
                )
            } else {
                faceResult
            }
        }

        val updatedValidImages = updatedFaceResults
            .filter { it.uri !in excludedImages } // Filter excluded
            .filter { it.hasFace }
            .filter { it.croppedFaceBitmap != null }
            .filter { it.errorMessage == null }
            .filter { it.faceCount >= 1 }
            .map { faceResult ->
                TrainingSanityChecker.ValidTrainingImage(
                    uri = faceResult.uri,
                    croppedFaceBitmap = faceResult.croppedFaceBitmap!!,
                    faceCount = faceResult.faceCount
                )
            }

        val updatedErrors = result.validationErrors.toMutableList()

        // Remove errors for manually selected faces or excluded images
        updatedErrors.removeAll { error ->
            when (error) {
                is TrainingSanityChecker.ValidationError.MultipleFacesDetected ->
                    manualFaceSelections.containsKey(error.uri) || excludedImages.contains(error.uri)
                is TrainingSanityChecker.ValidationError.NoFaceDetected ->
                    error.uris.any { excludedImages.contains(it) }
                is TrainingSanityChecker.ValidationError.ImageLoadError ->
                    excludedImages.contains(error.uri)
                else -> false
            }
        }

        // Update insufficient images error
        if (updatedValidImages.size < 15) {
            if (updatedErrors.none { it is TrainingSanityChecker.ValidationError.InsufficientImages }) {
                updatedErrors.add(
                    TrainingSanityChecker.ValidationError.InsufficientImages(
                        required = 15,
                        available = updatedValidImages.size
                    )
                )
            }
        } else {
            updatedErrors.removeAll { it is TrainingSanityChecker.ValidationError.InsufficientImages }
        }

        val isValid = updatedErrors.isEmpty() && updatedValidImages.size >= 15

        return result.copy(
            isValid = isValid,
            faceDetectionResults = updatedFaceResults,
            validationErrors = updatedErrors,
            validImagesWithFaces = updatedValidImages,
            excludedImages = excludedImages.toSet()  // Immutable copy for Compose state detection
        )
    }

    fun getFormattedErrors(result: TrainingSanityChecker.SanityCheckResult): List<String> {
        return sanityChecker.formatValidationErrors(result.validationErrors)
    }

    fun reset() {
        _uiState.value = ScanningState.Idle
        _trainingState.value = TrainingState.Idle
        currentImageUris = emptyList()
        manualFaceSelections.clear()
        excludedImages.clear()
        personInfo = null
    }

    override fun onCleared() {
        super.onCleared()
        sanityChecker.cleanup()
        faceDetectionHelper.cleanup()
        faceNetModel.close()
    }
}

// Extension functions for copying results
private fun FaceDetectionHelper.FaceDetectionResult.copy(
    uri: Uri = this.uri,
    hasFace: Boolean = this.hasFace,
    faceCount: Int = this.faceCount,
    faceBounds: List<android.graphics.Rect> = this.faceBounds,
    croppedFaceBitmap: Bitmap? = this.croppedFaceBitmap,
    errorMessage: String? = this.errorMessage
): FaceDetectionHelper.FaceDetectionResult {
    return FaceDetectionHelper.FaceDetectionResult(
        uri = uri,
        hasFace = hasFace,
        faceCount = faceCount,
        faceBounds = faceBounds,
        croppedFaceBitmap = croppedFaceBitmap,
        errorMessage = errorMessage
    )
}

private fun TrainingSanityChecker.SanityCheckResult.copy(
    isValid: Boolean = this.isValid,
    faceDetectionResults: List<FaceDetectionHelper.FaceDetectionResult> = this.faceDetectionResults,
    duplicateCheckResult: DuplicateImageDetector.DuplicateCheckResult = this.duplicateCheckResult,
    validationErrors: List<TrainingSanityChecker.ValidationError> = this.validationErrors,
    warnings: List<String> = this.warnings,
    validImagesWithFaces: List<TrainingSanityChecker.ValidTrainingImage> = this.validImagesWithFaces,
    excludedImages: Set<Uri> = this.excludedImages
): TrainingSanityChecker.SanityCheckResult {
    return TrainingSanityChecker.SanityCheckResult(
        isValid = isValid,
        faceDetectionResults = faceDetectionResults,
        duplicateCheckResult = duplicateCheckResult,
        validationErrors = validationErrors,
        warnings = warnings,
        validImagesWithFaces = validImagesWithFaces,
        excludedImages = excludedImages
    )
}