package com.placeholder.sherpai2.ui.trainingprep

import androidx.compose.animation.AnimatedVisibility
import androidx.compose.animation.core.animateFloatAsState
import androidx.compose.foundation.BorderStroke
import androidx.compose.foundation.ExperimentalFoundationApi
import androidx.compose.foundation.background
import androidx.compose.foundation.combinedClickable
import androidx.compose.foundation.layout.*
import androidx.compose.foundation.lazy.grid.*
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.alpha
import androidx.compose.ui.graphics.Color
import androidx.compose.ui.layout.ContentScale
import androidx.compose.ui.text.font.FontWeight
import androidx.compose.ui.unit.dp
import androidx.hilt.navigation.compose.hiltViewModel
import androidx.lifecycle.compose.collectAsStateWithLifecycle
import coil.compose.AsyncImage
import com.placeholder.sherpai2.data.local.entity.ImageEntity

/**
 * TrainingPhotoSelectorScreen - PREMIUM GRID + ROLLING SCAN
 *
 * FLOW:
 * 1. Shows PREMIUM faces only (solo, large, frontal)
 * 2. User picks 1-3 seed photos
 * 3. "Find Similar" button appears → launches RollingScanScreen
 * 4. Toggle to show all photos if needed
 */
@OptIn(ExperimentalMaterial3Api::class, ExperimentalFoundationApi::class)
@Composable
fun TrainingPhotoSelectorScreen(
    onBack: () -> Unit,
    onPhotosSelected: (List<android.net.Uri>) -> Unit,
    onLaunchRollingScan: ((List<String>) -> Unit)? = null,  // NEW: Navigate to rolling scan
    viewModel: TrainingPhotoSelectorViewModel = hiltViewModel()
) {
    val photos by viewModel.photosWithFaces.collectAsStateWithLifecycle()
    val selectedPhotos by viewModel.selectedPhotos.collectAsStateWithLifecycle()
    val isLoading by viewModel.isLoading.collectAsStateWithLifecycle()
    val isRanking by viewModel.isRanking.collectAsStateWithLifecycle()
    val showPremiumOnly by viewModel.showPremiumOnly.collectAsStateWithLifecycle()
    val premiumCount by viewModel.premiumCount.collectAsStateWithLifecycle()
    val embeddingProgress by viewModel.embeddingProgress.collectAsStateWithLifecycle()

    Scaffold(
        topBar = {
            TopAppBar(
                title = {
                    Column {
                        Row(
                            horizontalArrangement = Arrangement.spacedBy(8.dp),
                            verticalAlignment = Alignment.CenterVertically
                        ) {
                            Text(
                                if (selectedPhotos.isEmpty()) {
                                    "Select Training Photos"
                                } else {
                                    "${selectedPhotos.size} selected"
                                },
                                style = MaterialTheme.typography.titleLarge,
                                fontWeight = FontWeight.Bold
                            )

                            // NEW: Ranking indicator
                            if (isRanking) {
                                CircularProgressIndicator(
                                    modifier = Modifier.size(16.dp),
                                    strokeWidth = 2.dp,
                                    color = MaterialTheme.colorScheme.primary
                                )
                            } else if (selectedPhotos.isNotEmpty()) {
                                Icon(
                                    Icons.Default.AutoAwesome,
                                    contentDescription = "AI Ranked",
                                    modifier = Modifier.size(20.dp),
                                    tint = MaterialTheme.colorScheme.primary
                                )
                            }
                        }

                        // Status text
                        Text(
                            when {
                                isRanking -> "Ranking similar photos..."
                                showPremiumOnly -> "Showing $premiumCount premium faces"
                                else -> "Showing ${photos.size} photos with faces"
                            },
                            style = MaterialTheme.typography.bodySmall,
                            color = when {
                                isRanking -> MaterialTheme.colorScheme.primary
                                showPremiumOnly -> MaterialTheme.colorScheme.tertiary
                                else -> MaterialTheme.colorScheme.onSurfaceVariant
                            }
                        )
                    }
                },
                navigationIcon = {
                    IconButton(onClick = onBack) {
                        Icon(Icons.Default.ArrowBack, "Back")
                    }
                },
                actions = {
                    // Toggle premium/all
                    IconButton(onClick = { viewModel.togglePremiumOnly() }) {
                        Icon(
                            if (showPremiumOnly) Icons.Default.Star else Icons.Default.GridView,
                            contentDescription = if (showPremiumOnly) "Show all" else "Show premium only",
                            tint = if (showPremiumOnly) MaterialTheme.colorScheme.tertiary else MaterialTheme.colorScheme.onSurface
                        )
                    }
                    if (selectedPhotos.isNotEmpty()) {
                        TextButton(onClick = { viewModel.clearSelection() }) {
                            Text("Clear")
                        }
                    }
                },
                colors = TopAppBarDefaults.topAppBarColors(
                    containerColor = MaterialTheme.colorScheme.primaryContainer
                )
            )
        },
        bottomBar = {
            AnimatedVisibility(visible = selectedPhotos.isNotEmpty()) {
                SelectionBottomBar(
                    selectedCount = selectedPhotos.size,
                    canLaunchRollingScan = viewModel.canLaunchRollingScan && onLaunchRollingScan != null,
                    onClear = { viewModel.clearSelection() },
                    onFindSimilar = {
                        onLaunchRollingScan?.invoke(viewModel.getSeedImageIds())
                    },
                    onContinue = {
                        val uris = selectedPhotos.map { android.net.Uri.parse(it.imageUri) }
                        onPhotosSelected(uris)
                    }
                )
            }
        }
    ) { paddingValues ->
        Box(
            modifier = Modifier
                .fillMaxSize()
                .padding(paddingValues)
        ) {
            when {
                isLoading -> {
                    Box(
                        modifier = Modifier.fillMaxSize(),
                        contentAlignment = Alignment.Center
                    ) {
                        Column(
                            horizontalAlignment = Alignment.CenterHorizontally,
                            verticalArrangement = Arrangement.spacedBy(16.dp)
                        ) {
                            CircularProgressIndicator()
                            // Capture value to avoid race condition
                            val progress = embeddingProgress
                            if (progress != null) {
                                Text(
                                    "Preparing faces: ${progress.current}/${progress.total}",
                                    style = MaterialTheme.typography.bodyMedium,
                                    color = MaterialTheme.colorScheme.onSurfaceVariant
                                )
                                LinearProgressIndicator(
                                    progress = { progress.current.toFloat() / progress.total },
                                    modifier = Modifier
                                        .width(200.dp)
                                        .padding(top = 8.dp)
                                )
                            } else {
                                Text(
                                    "Loading premium faces...",
                                    style = MaterialTheme.typography.bodyMedium,
                                    color = MaterialTheme.colorScheme.onSurfaceVariant
                                )
                            }
                        }
                    }
                }
                photos.isEmpty() -> {
                    EmptyState(onBack)
                }
                else -> {
                    PhotoGrid(
                        photos = photos,
                        selectedPhotos = selectedPhotos,
                        onPhotoClick = { photo -> viewModel.toggleSelection(photo) }
                    )
                }
            }
        }
    }
}

@Composable
private fun SelectionBottomBar(
    selectedCount: Int,
    canLaunchRollingScan: Boolean,
    onClear: () -> Unit,
    onFindSimilar: () -> Unit,
    onContinue: () -> Unit
) {
    Surface(
        modifier = Modifier.fillMaxWidth(),
        color = MaterialTheme.colorScheme.primaryContainer,
        shadowElevation = 8.dp
    ) {
        Column(
            modifier = Modifier
                .fillMaxWidth()
                .padding(16.dp)
        ) {
            Row(
                modifier = Modifier.fillMaxWidth(),
                horizontalArrangement = Arrangement.SpaceBetween,
                verticalAlignment = Alignment.CenterVertically
            ) {
                Column {
                    Text(
                        "$selectedCount seed${if (selectedCount != 1) "s" else ""} selected",
                        style = MaterialTheme.typography.titleMedium,
                        fontWeight = FontWeight.Bold
                    )
                    Text(
                        when {
                            selectedCount == 0 -> "Pick 1-3 clear photos of the same person"
                            selectedCount in 1..3 -> "Tap 'Find Similar' to discover more"
                            selectedCount < 15 -> "Need ${15 - selectedCount} more for training"
                            else -> "Ready to train!"
                        },
                        style = MaterialTheme.typography.bodySmall,
                        color = when {
                            selectedCount in 1..3 -> MaterialTheme.colorScheme.tertiary
                            selectedCount < 15 -> MaterialTheme.colorScheme.error
                            else -> MaterialTheme.colorScheme.onPrimaryContainer.copy(alpha = 0.8f)
                        }
                    )
                }

                OutlinedButton(onClick = onClear) {
                    Text("Clear")
                }
            }

            Spacer(Modifier.height(12.dp))

            Row(
                modifier = Modifier.fillMaxWidth(),
                horizontalArrangement = Arrangement.spacedBy(8.dp)
            ) {
                // Find Similar button (prominent when 1-5 seeds selected)
                Button(
                    onClick = onFindSimilar,
                    enabled = canLaunchRollingScan,
                    modifier = Modifier.weight(1f),
                    colors = ButtonDefaults.buttonColors(
                        containerColor = MaterialTheme.colorScheme.tertiary
                    )
                ) {
                    Icon(
                        Icons.Default.AutoAwesome,
                        contentDescription = null,
                        modifier = Modifier.size(20.dp)
                    )
                    Spacer(Modifier.width(8.dp))
                    Text("Find Similar")
                }

                // Continue button (for manual selection path)
                Button(
                    onClick = onContinue,
                    enabled = selectedCount >= 15,
                    modifier = Modifier.weight(1f)
                ) {
                    Icon(
                        Icons.Default.Check,
                        contentDescription = null,
                        modifier = Modifier.size(20.dp)
                    )
                    Spacer(Modifier.width(8.dp))
                    Text("Train ($selectedCount)")
                }
            }
        }
    }
}

@OptIn(ExperimentalFoundationApi::class)
@Composable
private fun PhotoGrid(
    photos: List<ImageEntity>,
    selectedPhotos: Set<ImageEntity>,
    onPhotoClick: (ImageEntity) -> Unit
) {
    LazyVerticalGrid(
        columns = GridCells.Fixed(3),
        contentPadding = PaddingValues(
            start = 4.dp,
            end = 4.dp,
            bottom = 100.dp
        ),
        horizontalArrangement = Arrangement.spacedBy(4.dp),
        verticalArrangement = Arrangement.spacedBy(4.dp)
    ) {
        items(
            items = photos,
            key = { it.imageId }
        ) { photo ->
            PhotoThumbnail(
                photo = photo,
                isSelected = photo in selectedPhotos,
                onClick = { onPhotoClick(photo) }
            )
        }
    }
}

@OptIn(ExperimentalFoundationApi::class)
@Composable
private fun PhotoThumbnail(
    photo: ImageEntity,
    isSelected: Boolean,
    onClick: () -> Unit
) {
    // NEW: Fade animation for non-selected photos
    val alpha by animateFloatAsState(
        targetValue = if (isSelected) 1f else 1f,
        label = "photoAlpha"
    )

    Card(
        modifier = Modifier
            .fillMaxWidth()
            .aspectRatio(1f)
            .alpha(alpha)
            .combinedClickable(onClick = onClick),
        shape = RoundedCornerShape(4.dp),
        border = if (isSelected) {
            BorderStroke(4.dp, MaterialTheme.colorScheme.primary)
        } else null
    ) {
        Box {
            // Photo
            AsyncImage(
                model = photo.imageUri,
                contentDescription = null,
                modifier = Modifier.fillMaxSize(),
                contentScale = ContentScale.Crop
            )

            // Face count badge (top-left)
            if (photo.faceCount != null && photo.faceCount!! > 0) {
                Surface(
                    modifier = Modifier
                        .align(Alignment.TopStart)
                        .padding(4.dp),
                    shape = CircleShape,
                    color = MaterialTheme.colorScheme.primaryContainer.copy(alpha = 0.95f)
                ) {
                    Row(
                        modifier = Modifier.padding(horizontal = 6.dp, vertical = 2.dp),
                        horizontalArrangement = Arrangement.spacedBy(2.dp),
                        verticalAlignment = Alignment.CenterVertically
                    ) {
                        Icon(
                            Icons.Default.Face,
                            contentDescription = null,
                            modifier = Modifier.size(12.dp),
                            tint = MaterialTheme.colorScheme.onPrimaryContainer
                        )
                        Text(
                            "${photo.faceCount}",
                            style = MaterialTheme.typography.labelSmall,
                            color = MaterialTheme.colorScheme.onPrimaryContainer,
                            fontWeight = FontWeight.Bold
                        )
                    }
                }
            }

            // Selection checkmark (top-right)
            if (isSelected) {
                Surface(
                    modifier = Modifier
                        .align(Alignment.TopEnd)
                        .padding(4.dp)
                        .size(28.dp),
                    shape = CircleShape,
                    color = MaterialTheme.colorScheme.primary,
                    shadowElevation = 4.dp
                ) {
                    Box(contentAlignment = Alignment.Center) {
                        Icon(
                            Icons.Default.CheckCircle,
                            contentDescription = "Selected",
                            modifier = Modifier.size(20.dp),
                            tint = MaterialTheme.colorScheme.onPrimary
                        )
                    }
                }
            }

            // Dim overlay when selected
            if (isSelected) {
                Box(
                    modifier = Modifier
                        .fillMaxSize()
                        .background(Color.Black.copy(alpha = 0.2f))
                )
            }
        }
    }
}

@Composable
private fun EmptyState(onBack: () -> Unit) {
    Box(
        modifier = Modifier.fillMaxSize(),
        contentAlignment = Alignment.Center
    ) {
        Column(
            horizontalAlignment = Alignment.CenterHorizontally,
            verticalArrangement = Arrangement.spacedBy(16.dp),
            modifier = Modifier.padding(32.dp)
        ) {
            Icon(
                Icons.Default.SearchOff,
                contentDescription = null,
                modifier = Modifier.size(72.dp),
                tint = MaterialTheme.colorScheme.outline
            )

            Text(
                "No Photos with Faces Found",
                style = MaterialTheme.typography.titleLarge,
                fontWeight = FontWeight.Bold
            )

            Text(
                "Make sure the face detection cache has scanned your library",
                style = MaterialTheme.typography.bodyMedium,
                color = MaterialTheme.colorScheme.onSurfaceVariant
            )

            Button(onClick = onBack) {
                Icon(Icons.Default.ArrowBack, null)
                Spacer(Modifier.width(8.dp))
                Text("Go Back")
            }
        }
    }
}