package com.placeholder.sherpai2.ui.trainingprep

import android.net.Uri
import androidx.activity.compose.rememberLauncherForActivityResult
import androidx.activity.result.PickVisualMediaRequest
import androidx.activity.result.contract.ActivityResultContracts
import androidx.compose.foundation.BorderStroke
import androidx.compose.foundation.Image
import androidx.compose.foundation.background
import androidx.compose.foundation.border
import androidx.compose.foundation.layout.*
import androidx.compose.foundation.lazy.LazyColumn
import androidx.compose.foundation.lazy.itemsIndexed
import androidx.compose.foundation.shape.CircleShape
import androidx.compose.foundation.shape.RoundedCornerShape
import androidx.compose.material.icons.Icons
import androidx.compose.material.icons.filled.*
import androidx.compose.material3.*
import androidx.compose.runtime.*
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.draw.clip
import androidx.compose.ui.graphics.Color
import androidx.compose.ui.graphics.asImageBitmap
import androidx.compose.ui.layout.ContentScale
import androidx.compose.ui.text.font.FontWeight
import androidx.compose.ui.text.style.TextAlign
import androidx.compose.ui.unit.dp
import androidx.hilt.navigation.compose.hiltViewModel
import coil.compose.AsyncImage


@OptIn(ExperimentalMaterial3Api::class)
@Composable
fun ScanResultsScreen(
    state: ScanningState,
    onFinish: () -> Unit,
    trainViewModel: TrainViewModel = hiltViewModel()
) {
    var showFacePickerDialog by remember { mutableStateOf<FaceDetectionHelper.FaceDetectionResult?>(null) }
    val trainingState by trainViewModel.trainingState.collectAsState()

    LaunchedEffect(trainingState) {
        when (trainingState) {
            is TrainingState.Success -> {
                trainViewModel.resetTrainingState()
                onFinish()
            }
            is TrainingState.Error -> {}
            else -> {}
        }
    }

    // No Scaffold - MainScreen provides TopAppBar
    Box(modifier = Modifier.fillMaxSize()) {
        when (state) {
            is ScanningState.Idle -> {}

            is ScanningState.Processing -> {
                ProcessingView(progress = state.progress, total = state.total)
            }

            is ScanningState.Success -> {
                ImprovedResultsView(
                    result = state.sanityCheckResult,
                    onContinue = {
                        trainViewModel.createFaceModel(
                            trainViewModel.getPersonInfo()?.name ?: "Unknown"
                        )
                    },
                    onRetry = onFinish,
                    onReplaceImage = { oldUri, newUri ->
                        trainViewModel.replaceImage(oldUri, newUri)
                    },
                    onSelectFaceFromMultiple = { result ->
                        showFacePickerDialog = result
                    },
                    trainViewModel = trainViewModel
                )
            }

            is ScanningState.Error -> {
                ErrorView(message = state.message, onRetry = onFinish)
            }
        }

        if (trainingState is TrainingState.Processing) {
            TrainingOverlay(trainingState = trainingState as TrainingState.Processing)
        }
    }

    showFacePickerDialog?.let { result ->
        FacePickerDialog(
            result = result,
            onDismiss = { showFacePickerDialog = null },
            onFaceSelected = { faceIndex, croppedFaceBitmap ->
                trainViewModel.selectFaceFromImage(result.uri, faceIndex, croppedFaceBitmap)
                showFacePickerDialog = null
            }
        )
    }
}

@Composable
private fun TrainingOverlay(trainingState: TrainingState.Processing) {
    Box(
        modifier = Modifier.fillMaxSize().background(Color.Black.copy(alpha = 0.7f)),
        contentAlignment = Alignment.Center
    ) {
        Card(
            modifier = Modifier.padding(32.dp).fillMaxWidth(0.9f),
            colors = CardDefaults.cardColors(containerColor = MaterialTheme.colorScheme.surface)
        ) {
            Column(
                modifier = Modifier.padding(24.dp),
                horizontalAlignment = Alignment.CenterHorizontally,
                verticalArrangement = Arrangement.spacedBy(16.dp)
            ) {
                CircularProgressIndicator(modifier = Modifier.size(64.dp), strokeWidth = 6.dp)
                Text("Creating Face Model", style = MaterialTheme.typography.titleLarge, fontWeight = FontWeight.Bold)
                Text(trainingState.stage, style = MaterialTheme.typography.bodyMedium, textAlign = TextAlign.Center, color = MaterialTheme.colorScheme.onSurfaceVariant)
                if (trainingState.total > 0) {
                    LinearProgressIndicator(
                        progress = { (trainingState.progress.toFloat() / trainingState.total.toFloat()).coerceIn(0f, 1f) },
                        modifier = Modifier.fillMaxWidth()
                    )
                    Text("${trainingState.progress} / ${trainingState.total}", style = MaterialTheme.typography.bodySmall, color = MaterialTheme.colorScheme.onSurfaceVariant)
                }
            }
        }
    }
}

@Composable
private fun ProcessingView(progress: Int, total: Int) {
    Column(
        modifier = Modifier.fillMaxSize(),
        horizontalAlignment = Alignment.CenterHorizontally,
        verticalArrangement = Arrangement.Center
    ) {
        CircularProgressIndicator(modifier = Modifier.size(64.dp), strokeWidth = 6.dp)
        Spacer(modifier = Modifier.height(24.dp))
        Text("Analyzing images...", style = MaterialTheme.typography.titleMedium)
        Spacer(modifier = Modifier.height(8.dp))
        Text("Detecting faces and checking for duplicates", style = MaterialTheme.typography.bodyMedium, color = MaterialTheme.colorScheme.onSurfaceVariant)
        if (total > 0) {
            Spacer(modifier = Modifier.height(16.dp))
            LinearProgressIndicator(
                progress = { (progress.toFloat() / total.toFloat()).coerceIn(0f, 1f) },
                modifier = Modifier.width(200.dp)
            )
            Text("$progress / $total", style = MaterialTheme.typography.bodySmall)
        }
    }
}

@Composable
private fun ImprovedResultsView(
    result: TrainingSanityChecker.SanityCheckResult,
    onContinue: () -> Unit,
    onRetry: () -> Unit,
    onReplaceImage: (Uri, Uri) -> Unit,
    onSelectFaceFromMultiple: (FaceDetectionHelper.FaceDetectionResult) -> Unit,
    trainViewModel: TrainViewModel
) {
    LazyColumn(
        modifier = Modifier.fillMaxSize(),
        contentPadding = PaddingValues(16.dp),
        verticalArrangement = Arrangement.spacedBy(16.dp)
    ) {
        item {
            Card(
                modifier = Modifier.fillMaxWidth(),
                colors = CardDefaults.cardColors(containerColor = MaterialTheme.colorScheme.secondaryContainer)
            ) {
                Column(modifier = Modifier.padding(16.dp)) {
                    Text("Analysis Complete!", style = MaterialTheme.typography.headlineSmall, fontWeight = FontWeight.Bold)
                    Spacer(modifier = Modifier.height(4.dp))
                    Text(
                        "Review your images below. Tap 'Pick Face' on group photos to choose which person to train on, or 'Replace' to swap out any image.",
                        style = MaterialTheme.typography.bodyMedium,
                        color = MaterialTheme.colorScheme.onSecondaryContainer.copy(alpha = 0.8f)
                    )
                }
            }
        }

        item {
            ProgressSummaryCard(
                totalImages = result.faceDetectionResults.size,
                validImages = result.validImagesWithFaces.size,
                requiredImages = 10,
                isValid = result.isValid
            )
        }

        item {
            Text("Your Images (${result.faceDetectionResults.size})", style = MaterialTheme.typography.titleLarge, fontWeight = FontWeight.Bold)
        }

        itemsIndexed(result.faceDetectionResults) { index, imageResult ->
            ImageResultCard(
                index = index + 1,
                result = imageResult,
                onReplace = { newUri -> onReplaceImage(imageResult.uri, newUri) },
                onSelectFace = if (imageResult.faceCount > 1) { { onSelectFaceFromMultiple(imageResult) } } else null,
                trainViewModel = trainViewModel,
                isExcluded = trainViewModel.isImageExcluded(imageResult.uri)
            )
        }

        if (result.validationErrors.isNotEmpty()) {
            item {
                Spacer(modifier = Modifier.height(8.dp))
                ValidationIssuesCard(errors = result.validationErrors, trainViewModel = trainViewModel)
            }
        }

        item {
            Spacer(modifier = Modifier.height(8.dp))
            Button(
                onClick = if (result.isValid) onContinue else onRetry,
                modifier = Modifier.fillMaxWidth(),
                enabled = result.isValid,
                colors = ButtonDefaults.buttonColors(
                    containerColor = if (result.isValid) MaterialTheme.colorScheme.primary else MaterialTheme.colorScheme.error.copy(alpha = 0.5f)
                )
            ) {
                Icon(if (result.isValid) Icons.Default.CheckCircle else Icons.Default.Warning, contentDescription = null)
                Spacer(modifier = Modifier.width(8.dp))
                Text(
                    if (result.isValid)
                        "Continue to Training (${result.validImagesWithFaces.size} images)"
                    else
                        "Fix ${result.validationErrors.size} Issue${if (result.validationErrors.size != 1) "s" else ""} to Continue"
                )
            }

            if (!result.isValid) {
                Spacer(modifier = Modifier.height(8.dp))
                Surface(
                    modifier = Modifier.fillMaxWidth(),
                    color = MaterialTheme.colorScheme.tertiaryContainer,
                    shape = RoundedCornerShape(8.dp)
                ) {
                    Row(modifier = Modifier.padding(12.dp), verticalAlignment = Alignment.CenterVertically) {
                        Icon(Icons.Default.Info, contentDescription = null, tint = MaterialTheme.colorScheme.onTertiaryContainer, modifier = Modifier.size(20.dp))
                        Spacer(modifier = Modifier.width(8.dp))
                        Text(
                            "Tip: Use 'Replace' to swap problematic images, or 'Pick Face' to choose from group photos",
                            style = MaterialTheme.typography.bodySmall,
                            color = MaterialTheme.colorScheme.onTertiaryContainer
                        )
                    }
                }
            }
        }
    }
}

@Composable
private fun ProgressSummaryCard(totalImages: Int, validImages: Int, requiredImages: Int, isValid: Boolean) {
    Card(
        modifier = Modifier.fillMaxWidth(),
        colors = CardDefaults.cardColors(
            containerColor = if (isValid) MaterialTheme.colorScheme.primaryContainer.copy(alpha = 0.5f) else MaterialTheme.colorScheme.errorContainer.copy(alpha = 0.3f)
        )
    ) {
        Column(modifier = Modifier.padding(16.dp)) {
            Row(modifier = Modifier.fillMaxWidth(), horizontalArrangement = Arrangement.SpaceBetween, verticalAlignment = Alignment.CenterVertically) {
                Text("Progress", style = MaterialTheme.typography.titleMedium, fontWeight = FontWeight.Bold)
                Icon(
                    imageVector = if (isValid) Icons.Default.CheckCircle else Icons.Default.Warning,
                    contentDescription = null,
                    tint = if (isValid) MaterialTheme.colorScheme.primary else MaterialTheme.colorScheme.error,
                    modifier = Modifier.size(32.dp)
                )
            }
            Spacer(modifier = Modifier.height(12.dp))
            Row(modifier = Modifier.fillMaxWidth(), horizontalArrangement = Arrangement.SpaceEvenly) {
                StatItem("Total", totalImages.toString(), MaterialTheme.colorScheme.onSurface)
                StatItem("Valid", validImages.toString(), if (validImages >= requiredImages) MaterialTheme.colorScheme.primary else MaterialTheme.colorScheme.error)
                StatItem("Need", requiredImages.toString(), MaterialTheme.colorScheme.onSurface.copy(alpha = 0.6f))
            }
            Spacer(modifier = Modifier.height(12.dp))
            LinearProgressIndicator(
                progress = { (validImages.toFloat() / requiredImages.toFloat()).coerceIn(0f, 1f) },
                modifier = Modifier.fillMaxWidth(),
                color = if (isValid) MaterialTheme.colorScheme.primary else MaterialTheme.colorScheme.error
            )
        }
    }
}

@Composable
private fun StatItem(label: String, value: String, color: Color) {
    Column(horizontalAlignment = Alignment.CenterHorizontally) {
        Text(value, style = MaterialTheme.typography.headlineMedium, fontWeight = FontWeight.Bold, color = color)
        Text(label, style = MaterialTheme.typography.bodySmall, color = color.copy(alpha = 0.7f))
    }
}

@Composable
private fun ImageResultCard(
    index: Int,
    result: FaceDetectionHelper.FaceDetectionResult,
    onReplace: (Uri) -> Unit,
    onSelectFace: (() -> Unit)?,
    trainViewModel: TrainViewModel,
    isExcluded: Boolean
) {
    val photoPickerLauncher = rememberLauncherForActivityResult(contract = ActivityResultContracts.PickVisualMedia()) { uri -> uri?.let { onReplace(it) } }

    val status = when {
        isExcluded -> ImageStatus.EXCLUDED
        result.errorMessage != null -> ImageStatus.ERROR
        !result.hasFace -> ImageStatus.NO_FACE
        result.faceCount > 1 -> ImageStatus.MULTIPLE_FACES
        result.faceCount == 1 -> ImageStatus.VALID
        else -> ImageStatus.ERROR
    }

    Card(
        modifier = Modifier.fillMaxWidth(),
        colors = CardDefaults.cardColors(
            containerColor = when (status) {
                ImageStatus.VALID -> MaterialTheme.colorScheme.primaryContainer.copy(alpha = 0.3f)
                ImageStatus.MULTIPLE_FACES -> MaterialTheme.colorScheme.tertiaryContainer.copy(alpha = 0.4f)
                ImageStatus.EXCLUDED -> MaterialTheme.colorScheme.surfaceVariant.copy(alpha = 0.5f)
                else -> MaterialTheme.colorScheme.errorContainer.copy(alpha = 0.3f)
            }
        )
    ) {
        Row(modifier = Modifier.fillMaxWidth().padding(12.dp), verticalAlignment = Alignment.CenterVertically, horizontalArrangement = Arrangement.spacedBy(12.dp)) {
            Box(
                modifier = Modifier.size(40.dp).background(
                    color = when (status) {
                        ImageStatus.VALID -> MaterialTheme.colorScheme.primary
                        ImageStatus.MULTIPLE_FACES -> MaterialTheme.colorScheme.tertiary
                        ImageStatus.EXCLUDED -> MaterialTheme.colorScheme.outline
                        else -> MaterialTheme.colorScheme.error
                    },
                    shape = CircleShape
                ),
                contentAlignment = Alignment.Center
            ) {
                Text(index.toString(), style = MaterialTheme.typography.titleMedium, fontWeight = FontWeight.Bold, color = Color.White)
            }

            if (result.croppedFaceBitmap != null) {
                Image(
                    bitmap = result.croppedFaceBitmap.asImageBitmap(),
                    contentDescription = "Face",
                    modifier = Modifier.size(64.dp).clip(RoundedCornerShape(8.dp)).border(
                        BorderStroke(2.dp, when (status) {
                            ImageStatus.VALID -> MaterialTheme.colorScheme.primary
                            ImageStatus.MULTIPLE_FACES -> MaterialTheme.colorScheme.tertiary
                            ImageStatus.EXCLUDED -> MaterialTheme.colorScheme.outline
                            else -> MaterialTheme.colorScheme.error
                        }),
                        RoundedCornerShape(8.dp)
                    ),
                    contentScale = ContentScale.Crop
                )
            } else {
                AsyncImage(model = result.uri, contentDescription = "Original image", modifier = Modifier.size(64.dp).clip(RoundedCornerShape(8.dp)), contentScale = ContentScale.Crop)
            }

            Column(modifier = Modifier.weight(1f)) {
                Row(verticalAlignment = Alignment.CenterVertically) {
                    Icon(
                        imageVector = when (status) {
                            ImageStatus.VALID -> Icons.Default.CheckCircle
                            ImageStatus.MULTIPLE_FACES -> Icons.Default.Info
                            ImageStatus.EXCLUDED -> Icons.Default.RemoveCircle
                            else -> Icons.Default.Warning
                        },
                        contentDescription = null,
                        tint = when (status) {
                            ImageStatus.VALID -> MaterialTheme.colorScheme.primary
                            ImageStatus.MULTIPLE_FACES -> MaterialTheme.colorScheme.tertiary
                            ImageStatus.EXCLUDED -> MaterialTheme.colorScheme.outline
                            else -> MaterialTheme.colorScheme.error
                        },
                        modifier = Modifier.size(20.dp)
                    )
                    Spacer(modifier = Modifier.width(4.dp))
                    Text(
                        text = when (status) {
                            ImageStatus.VALID -> "Face Detected"
                            ImageStatus.MULTIPLE_FACES -> "Multiple Faces (${result.faceCount})"
                            ImageStatus.NO_FACE -> "No Face Detected"
                            ImageStatus.EXCLUDED -> "Excluded"
                            ImageStatus.ERROR -> "Error"
                        },
                        style = MaterialTheme.typography.bodyMedium,
                        fontWeight = FontWeight.SemiBold
                    )
                }
                Text(result.uri.lastPathSegment ?: "Unknown", style = MaterialTheme.typography.bodySmall, color = MaterialTheme.colorScheme.onSurfaceVariant, maxLines = 1)
            }

            Column(horizontalAlignment = Alignment.End, verticalArrangement = Arrangement.spacedBy(4.dp)) {
                if (onSelectFace != null && !isExcluded) {
                    OutlinedButton(
                        onClick = onSelectFace,
                        modifier = Modifier.height(32.dp),
                        contentPadding = PaddingValues(horizontal = 12.dp, vertical = 0.dp),
                        colors = ButtonDefaults.outlinedButtonColors(contentColor = MaterialTheme.colorScheme.tertiary),
                        border = BorderStroke(1.dp, MaterialTheme.colorScheme.tertiary)
                    ) {
                        Icon(Icons.Default.Face, contentDescription = null, modifier = Modifier.size(16.dp))
                        Spacer(modifier = Modifier.width(4.dp))
                        Text("Pick Face", style = MaterialTheme.typography.bodySmall)
                    }
                }

                if (!isExcluded) {
                    OutlinedButton(
                        onClick = { photoPickerLauncher.launch(PickVisualMediaRequest(ActivityResultContracts.PickVisualMedia.ImageOnly)) },
                        modifier = Modifier.height(32.dp),
                        contentPadding = PaddingValues(horizontal = 12.dp, vertical = 0.dp)
                    ) {
                        Icon(Icons.Default.Refresh, contentDescription = null, modifier = Modifier.size(16.dp))
                        Spacer(modifier = Modifier.width(4.dp))
                        Text("Replace", style = MaterialTheme.typography.bodySmall)
                    }
                }

                OutlinedButton(
                    onClick = {
                        if (isExcluded) trainViewModel.includeImage(result.uri) else trainViewModel.excludeImage(result.uri)
                    },
                    modifier = Modifier.height(32.dp),
                    contentPadding = PaddingValues(horizontal = 12.dp, vertical = 0.dp),
                    colors = ButtonDefaults.outlinedButtonColors(contentColor = if (isExcluded) MaterialTheme.colorScheme.primary else MaterialTheme.colorScheme.error),
                    border = BorderStroke(1.dp, if (isExcluded) MaterialTheme.colorScheme.primary else MaterialTheme.colorScheme.error)
                ) {
                    Icon(if (isExcluded) Icons.Default.Add else Icons.Default.Close, contentDescription = null, modifier = Modifier.size(16.dp))
                    Spacer(modifier = Modifier.width(4.dp))
                    Text(if (isExcluded) "Include" else "Exclude", style = MaterialTheme.typography.bodySmall)
                }
            }
        }
    }
}

@Composable
private fun ValidationIssuesCard(errors: List<TrainingSanityChecker.ValidationError>, trainViewModel: TrainViewModel) {
    Card(
        modifier = Modifier.fillMaxWidth(),
        colors = CardDefaults.cardColors(containerColor = MaterialTheme.colorScheme.errorContainer.copy(alpha = 0.3f))
    ) {
        Column(modifier = Modifier.padding(16.dp), verticalArrangement = Arrangement.spacedBy(8.dp)) {
            Row(verticalAlignment = Alignment.CenterVertically) {
                Icon(Icons.Default.Warning, contentDescription = null, tint = MaterialTheme.colorScheme.error)
                Spacer(modifier = Modifier.width(8.dp))
                Text("Issues Found (${errors.size})", style = MaterialTheme.typography.titleMedium, fontWeight = FontWeight.Bold, color = MaterialTheme.colorScheme.error)
            }

            HorizontalDivider(color = MaterialTheme.colorScheme.error.copy(alpha = 0.3f))

            errors.forEach { error ->
                when (error) {
                    is TrainingSanityChecker.ValidationError.NoFaceDetected -> {
                        Text("• ${error.uris.size} image(s) without detected faces - use Replace button", style = MaterialTheme.typography.bodyMedium)
                    }
                    is TrainingSanityChecker.ValidationError.MultipleFacesDetected -> {
                        Text("• ${error.uri.lastPathSegment} has ${error.faceCount} faces - use Pick Face button", style = MaterialTheme.typography.bodyMedium)
                    }
                    is TrainingSanityChecker.ValidationError.DuplicateImages -> {
                        Column(verticalArrangement = Arrangement.spacedBy(8.dp)) {
                            Row(modifier = Modifier.fillMaxWidth(), horizontalArrangement = Arrangement.SpaceBetween, verticalAlignment = Alignment.CenterVertically) {
                                Text("• ${error.groups.size} duplicate group(s) found", style = MaterialTheme.typography.bodyMedium, modifier = Modifier.weight(1f))

                                Button(
                                    onClick = {
                                        error.groups.forEach { group ->
                                            group.images.drop(1).forEach { uri ->
                                                trainViewModel.excludeImage(uri)
                                            }
                                        }
                                    },
                                    colors = ButtonDefaults.buttonColors(containerColor = MaterialTheme.colorScheme.tertiary),
                                    modifier = Modifier.height(36.dp)
                                ) {
                                    Icon(Icons.Default.DeleteSweep, contentDescription = null, modifier = Modifier.size(16.dp))
                                    Spacer(Modifier.width(4.dp))
                                    Text("Drop All", style = MaterialTheme.typography.labelMedium)
                                }
                            }

                            Text("${error.groups.sumOf { it.images.size - 1 }} duplicate images will be excluded", style = MaterialTheme.typography.bodySmall, color = MaterialTheme.colorScheme.onSurface.copy(alpha = 0.6f))
                        }
                    }
                    is TrainingSanityChecker.ValidationError.InsufficientImages -> {
                        Text("• Need ${error.required} valid images, currently have ${error.available}", style = MaterialTheme.typography.bodyMedium, fontWeight = FontWeight.Bold)
                    }
                    is TrainingSanityChecker.ValidationError.ImageLoadError -> {
                        Text("• Failed to load ${error.uri.lastPathSegment} - use Replace button", style = MaterialTheme.typography.bodyMedium)
                    }
                }
            }
        }
    }
}

@Composable
private fun ErrorView(message: String, onRetry: () -> Unit) {
    Column(modifier = Modifier.fillMaxSize().padding(16.dp), horizontalAlignment = Alignment.CenterHorizontally, verticalArrangement = Arrangement.Center) {
        Icon(imageVector = Icons.Default.Close, contentDescription = null, modifier = Modifier.size(64.dp), tint = MaterialTheme.colorScheme.error)
        Spacer(modifier = Modifier.height(16.dp))
        Text("Error", style = MaterialTheme.typography.titleLarge, fontWeight = FontWeight.Bold)
        Spacer(modifier = Modifier.height(8.dp))
        Text(message, style = MaterialTheme.typography.bodyMedium, textAlign = TextAlign.Center)
        Spacer(modifier = Modifier.height(24.dp))
        Button(onClick = onRetry) {
            Icon(Icons.Default.Refresh, contentDescription = null)
            Spacer(modifier = Modifier.width(8.dp))
            Text("Try Again")
        }
    }
}

private enum class ImageStatus {
    VALID,
    MULTIPLE_FACES,
    NO_FACE,
    ERROR,
    EXCLUDED
}