package com.example.livingai.pages.camera

import android.graphics.Bitmap
import android.graphics.BitmapFactory
import android.net.Uri
import androidx.compose.foundation.Canvas
import androidx.compose.foundation.Image
import androidx.compose.foundation.layout.Arrangement
import androidx.compose.foundation.layout.Box
import androidx.compose.foundation.layout.Row
import androidx.compose.foundation.layout.fillMaxSize
import androidx.compose.foundation.layout.fillMaxWidth
import androidx.compose.foundation.layout.padding
import androidx.compose.material3.Button
import androidx.compose.material3.CircularProgressIndicator
import androidx.compose.material3.OutlinedButton
import androidx.compose.material3.Scaffold
import androidx.compose.material3.Text
import androidx.compose.runtime.Composable
import androidx.compose.runtime.LaunchedEffect
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.rememberCoroutineScope
import androidx.compose.runtime.setValue
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.geometry.Offset
import androidx.compose.ui.geometry.Size
import androidx.compose.ui.graphics.Color
import androidx.compose.ui.graphics.drawscope.Stroke
import androidx.compose.ui.graphics.toArgb
import androidx.compose.ui.layout.ContentScale
import androidx.compose.ui.platform.LocalContext
import androidx.compose.ui.unit.dp
import androidx.exifinterface.media.ExifInterface
import coil.compose.rememberAsyncImagePainter
import com.example.livingai.domain.ml.SubjectSegmenterHelper
import com.example.livingai.ui.theme.LivingAITheme
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext
import kotlin.math.min

@Composable
fun ViewImageScreen(
    imageUri: String,
    shouldAllowRetake: Boolean,
    showAccept: Boolean,
    showBack: Boolean,
    showSegment: Boolean = false,
    animalId: String,
    orientation: String? = null,
    onRetake: () -> Unit,
    onAccept: (String) -> Unit,
    onSegmented: (String) -> Unit = {},
    onBack: () -> Unit
) {
    val context = LocalContext.current
    val scope = rememberCoroutineScope()
    var boundingBox by remember { mutableStateOf<String?>(null) }
    var imageWidth by remember { mutableStateOf(0f) }
    var imageHeight by remember { mutableStateOf(0f) }

    val displayedUri = Uri.parse(imageUri) 
    var isSegmenting by remember { mutableStateOf(false) }

    val segmenterHelper = remember { SubjectSegmenterHelper(context) }

    LaunchedEffect(displayedUri) {
        val uri = displayedUri
        try {
            context.contentResolver.openInputStream(uri)?.use { inputStream ->
                val exif = ExifInterface(inputStream)
                boundingBox = exif.getAttribute(ExifInterface.TAG_USER_COMMENT)

                val width = exif.getAttributeInt(ExifInterface.TAG_IMAGE_WIDTH, 0)
                val height = exif.getAttributeInt(ExifInterface.TAG_IMAGE_LENGTH, 0)
                if (width > 0 && height > 0) {
                     imageWidth = width.toFloat()
                     imageHeight = height.toFloat()
                }
            }
        } catch (e: Exception) {
            e.printStackTrace()
        }
    }

    LivingAITheme {
        Scaffold {
            Box(modifier = Modifier
                .fillMaxSize()
                .padding(it)) {
                
                if (isSegmenting) {
                    CircularProgressIndicator(modifier = Modifier.align(Alignment.Center))
                } else {
                    Image(
                        painter = rememberAsyncImagePainter(
                            model = displayedUri,
                            onSuccess = { state ->
                                if (imageWidth == 0f || imageHeight == 0f) {
                                    imageWidth = state.result.drawable.intrinsicWidth.toFloat()
                                    imageHeight = state.result.drawable.intrinsicHeight.toFloat()
                                }
                            }
                        ),
                        contentDescription = "Captured Image",
                        modifier = Modifier.fillMaxSize(),
                        contentScale = ContentScale.Fit,
                        alignment = Alignment.Center
                    )

                    if (boundingBox != null && imageWidth > 0 && imageHeight > 0) {
                         Canvas(modifier = Modifier.fillMaxSize()) {
                             val canvasWidth = size.width
                             val canvasHeight = size.height

                             val parts = boundingBox!!.split(",")
                             if (parts.size == 4) {
                                 val left = parts[0].toFloatOrNull() ?: 0f
                                 val top = parts[1].toFloatOrNull() ?: 0f
                                 val right = parts[2].toFloatOrNull() ?: 0f
                                 val bottom = parts[3].toFloatOrNull() ?: 0f

                                 val widthRatio = canvasWidth / imageWidth
                                 val heightRatio = canvasHeight / imageHeight
                                 val scale = min(widthRatio, heightRatio)

                                 val displayedWidth = imageWidth * scale
                                 val displayedHeight = imageHeight * scale
                                 
                                 val offsetX = (canvasWidth - displayedWidth) / 2
                                 val offsetY = (canvasHeight - displayedHeight) / 2

                                 val rectLeft = left * scale + offsetX
                                 val rectTop = top * scale + offsetY
                                 val rectRight = right * scale + offsetX
                                 val rectBottom = bottom * scale + offsetY
                                 
                                 drawRect(
                                     color = Color.Yellow,
                                     topLeft = Offset(rectLeft, rectTop),
                                     size = Size(rectRight - rectLeft, rectBottom - rectTop),
                                     style = Stroke(width = 3.dp.toPx())
                                 )
                             }
                         }
                    }
                }

                Row(
                    modifier = Modifier
                        .fillMaxWidth()
                        .align(Alignment.BottomCenter)
                        .padding(16.dp),
                    horizontalArrangement = Arrangement.SpaceEvenly
                ) {
                    if (shouldAllowRetake && !isSegmenting) {
                        OutlinedButton(onClick = onRetake) {
                            Text("Retake")
                        }
                    }
                    
                    if (!isSegmenting) {
                        if (showSegment) {
                            Button(onClick = {
                                scope.launch {
                                    isSegmenting = true

                                    try {
                                        val originalBitmap = withContext(Dispatchers.IO) {
                                            context.contentResolver.openInputStream(displayedUri)?.use {
                                                BitmapFactory.decodeStream(it)
                                            }
                                        }

                                        if (originalBitmap == null) {
                                            isSegmenting = false
                                            return@launch
                                        }

                                        val finalBitmap = if (boundingBox != null) {

                                            val parts = boundingBox!!.split(",")
                                            if (parts.size == 4) {
                                                val left = parts[0].toFloatOrNull()?.toInt() ?: 0
                                                val top = parts[1].toFloatOrNull()?.toInt() ?: 0
                                                val right = parts[2].toFloatOrNull()?.toInt() ?: 0
                                                val bottom = parts[3].toFloatOrNull()?.toInt() ?: 0

                                                val w = right - left
                                                val h = bottom - top

                                                if (
                                                    w > 0 && h > 0 &&
                                                    left >= 0 && top >= 0 &&
                                                    left + w <= originalBitmap.width &&
                                                    top + h <= originalBitmap.height
                                                ) {
                                                    val cropped = Bitmap.createBitmap(originalBitmap, left, top, w, h)

                                                    val segmentedCrop =
                                                        segmenterHelper.segmentToBitmap(cropped)

                                                    if (segmentedCrop != null) {
                                                        Bitmap.createBitmap(
                                                            originalBitmap.width,
                                                            originalBitmap.height,
                                                            Bitmap.Config.ARGB_8888
                                                        ).apply {
                                                            android.graphics.Canvas(this).apply {
                                                                drawColor(Color.Black.toArgb())
                                                                drawBitmap(segmentedCrop, left.toFloat(), top.toFloat(), null)
                                                            }
                                                        }
                                                    } else {
                                                        null
                                                    }
                                                } else {
                                                    segmenterHelper.segmentToBitmap(originalBitmap)
                                                }
                                            } else {
                                                segmenterHelper.segmentToBitmap(originalBitmap)
                                            }
                                        } else {
                                            segmenterHelper.segmentToBitmap(originalBitmap)
                                        }

                                        val resultUri = finalBitmap?.let {
                                            segmenterHelper.segmentAndSave(
                                                it,
                                                animalId,
                                                orientation ?: "unknown",
                                                "Segmented images"
                                            )
                                        }

                                        if (resultUri != null) {
                                            onSegmented(resultUri.toString())
                                        }

                                    } catch (e: Exception) {
                                        e.printStackTrace()
                                    } finally {
                                        isSegmenting = false
                                    }
                                }
                            }) {
                                Text("Segment")
                            }
                        }
                        
                        if (showAccept) {
                            Button(onClick = { onAccept(displayedUri.toString()) }) {
                                Text("Accept")
                            }
                        }
                        
                        if (showBack || !showAccept) {
                            Button(onClick = onBack) {
                                Text("Back")
                            }
                        }
                    }
                }
            }
        }
    }
}
