package com.example.myapplication

import ImageCaptioner
import ai.onnxruntime.OrtEnvironment
import ai.onnxruntime.OrtSession
import android.content.Context
import android.content.Intent
import android.graphics.Bitmap
import android.graphics.BitmapFactory
import android.os.Bundle
import android.util.Log
import androidx.activity.ComponentActivity
import androidx.activity.compose.setContent
import androidx.activity.enableEdgeToEdge
import androidx.activity.result.contract.ActivityResultContracts
import androidx.compose.foundation.layout.fillMaxSize
import androidx.compose.foundation.layout.padding
import androidx.compose.material3.Scaffold
import androidx.compose.material3.Text
import androidx.compose.ui.Modifier
import com.example.myapplication.ui.theme.MyApplicationTheme
import android.graphics.ImageDecoder
import android.media.Image
import android.net.Uri
import android.os.Build
import android.provider.MediaStore
import android.view.View
import android.widget.Toast
import androidx.activity.compose.rememberLauncherForActivityResult
import androidx.annotation.RequiresApi
import androidx.camera.core.CameraSelector
import androidx.camera.core.ImageAnalysis
import androidx.camera.core.ImageCapture
import androidx.camera.core.ImageCaptureException
import androidx.camera.view.LifecycleCameraController
import androidx.camera.view.PreviewView
import androidx.compose.foundation.Image
import androidx.compose.foundation.layout.Arrangement
import androidx.compose.foundation.layout.Column
import androidx.compose.foundation.layout.Spacer
import androidx.compose.foundation.layout.height
import androidx.compose.material3.Button
import androidx.compose.runtime.Composable
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.setValue
import androidx.compose.ui.Alignment
import androidx.compose.ui.graphics.asImageBitmap
import androidx.compose.ui.graphics.decodeToImageBitmap
import androidx.compose.ui.platform.LocalContext
import androidx.compose.ui.platform.LocalLifecycleOwner
import androidx.compose.ui.unit.dp
import androidx.compose.ui.viewinterop.AndroidView
import androidx.lifecycle.lifecycleScope
import com.example.blipcaptioning.BLIPInference
import com.example.myapplication.SimpleTokenizer.Companion
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.asCoroutineDispatcher
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext
import java.io.BufferedReader
import java.io.ByteArrayOutputStream
import java.io.File
import java.io.InputStream
import java.io.InputStreamReader
import java.util.concurrent.ExecutorService
import java.util.concurrent.Executors
import java.util.concurrent.atomic.AtomicBoolean


class MainActivity : ComponentActivity() {

    private lateinit var context: Context
    private lateinit var modelFile: File
    private val TAG = "TAG"
    private lateinit var imageCaptioner: ImageCaptioner
    private lateinit var cameraExecutor: ExecutorService
    private val processing = AtomicBoolean(false)


    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)
        enableEdgeToEdge()
        context = applicationContext
        cameraExecutor = Executors.newFixedThreadPool(1)
        val cameraController = LifecycleCameraController(this)
        cameraController.bindToLifecycle(this)
        cameraController.cameraSelector = CameraSelector.DEFAULT_BACK_CAMERA
        initialize()
        cameraController.setImageAnalysisBackpressureStrategy(ImageAnalysis.STRATEGY_KEEP_ONLY_LATEST)
        cameraController.setImageAnalysisAnalyzer(
            cameraExecutor
        ) { imageProxy ->
            Log.d("debug", "trying to enter at ${System.currentTimeMillis()}")
            if (!processing.compareAndSet(false,true)) {
                Log.d("Debug", "processing.......")
                imageProxy.close()
                return@setImageAnalysisAnalyzer
            }
            Log.d("debug", processing.toString())
            var bitmap = imageProxy.toBitmap()
            val maxDimension = 1024
            bitmap  = resizeBitmap(bitmap,
                minOf(bitmap.width, maxDimension),
                minOf(bitmap.height, maxDimension))
            imageProxy.close()
            val bef = System.currentTimeMillis()
            Log.d("debug", "Entered at $bef")
            val caption = imageCaptioner.generateCaption(bitmap)
            val time = System.currentTimeMillis() - bef
            Log.d("time", (time).toString())
            runOnUiThread {
                Toast.makeText(this@MainActivity, "$caption $time", Toast.LENGTH_SHORT).show()
            }
            Log.d("Generated Caption", caption)
            Log.d("debug", processing.toString())
            processing.set(false)
            Log.d("debug", processing.toString())
            Log.d("debug", "Exited at ${System.currentTimeMillis()}")


        }

        setContent {
            MyApplicationTheme {
                Scaffold(modifier = Modifier.fillMaxSize()) { innerPadding ->
                    Column(
                        modifier = Modifier
                            .fillMaxSize()
                            .padding(innerPadding),
                        horizontalAlignment = Alignment.CenterHorizontally,
                        verticalArrangement = Arrangement.Center
                    ) {
                        Button(
                            onClick = {
                                galleryLauncher.launch("image/*")
                            }
                        ) {
                            Text("Choose Image")
                        }
                        CameraPreview(cameraController)
                    }
                }
            }
        }

    }

    fun initialize() {

        var reverseVocab: MutableMap<Int, String> = mutableMapOf()
        val inputStream = context.resources.openRawResource(R.raw.blip_vocab)
        val reader = BufferedReader(InputStreamReader(inputStream))

        var tokenId = 0
        reader.forEachLine { line ->
            val trimmedLine = line.trim()
            if( trimmedLine.isNotEmpty() ){
                when{
                    trimmedLine.contains('\t') -> {
                        val parts = trimmedLine.split('\t', limit = 2)
                        if(parts.size ==2 ){
                            val id = parts[0].toIntOrNull()
                            val token = parts[1]
                            if( id!=null) {
                                reverseVocab[id] = token
                            }
                        }
                    }
                    trimmedLine.contains(' ') -> {
                        val parts = trimmedLine.split(' ', limit = 2 )
                        val firstPartAsInt = parts[0].toIntOrNull()
                        if( firstPartAsInt != null && parts.size == 2 ){
                            val id = firstPartAsInt
                            val token = parts[1]
                            reverseVocab[id] = token
                        }
                        else{
//                            reverseVocab[id] = token
                            tokenId++
                        }
                    }
                    else -> {
//                        reverseVocab[id] = token
                        tokenId++
                    }
                }
            }

        }

        reader.close()
        Log.d(TAG, "Vocabulary loaded:  tokens")

        val ortEnvironment = OrtEnvironment.getEnvironment()

        // Load vision encoder
        val visionModelBytes = context.resources.openRawResource(R.raw.blip_vision_encoder_quantized).readBytes()
        val visionSession = ortEnvironment!!.createSession(visionModelBytes)
        Log.d(TAG, "Vision encoder loaded successfully")

        // Load text decoder
        val textModelBytes = context.resources.openRawResource(R.raw.blip_text_decoder_quantized).readBytes()
        val textSession = ortEnvironment.createSession(textModelBytes)
        Log.d(TAG, "Text decoder loaded successfully")

        imageCaptioner = ImageCaptioner(
            ortEnvironment,
            visionSession,
            textSession,
            reverseVocab
        )
    }

    private fun loadBitmapFromUri(uri: Uri): Bitmap? {
        return try {
            val inputStream = contentResolver.openInputStream(uri)
            val bitmap = BitmapFactory.decodeStream(inputStream)
            inputStream?.close()

            // Fix: Use min instead of modulo operator
            val maxDimension = 1024
            resizeBitmap(bitmap,
                minOf(bitmap.width, maxDimension),
                minOf(bitmap.height, maxDimension))
        } catch (e: Exception) {
            Log.e(TAG, "Error loading bitmap from URI", e)
            null
        }
    }

    private fun resizeBitmap(bitmap: Bitmap, maxWidth: Int, maxHeight: Int): Bitmap {
        val width = bitmap.width
        val height = bitmap.height

        if (width <= maxWidth && height <= maxHeight) {
            return bitmap
        }

        val scaleWidth = maxWidth.toFloat() / width
        val scaleHeight = maxHeight.toFloat() / height
        val scale = minOf(scaleWidth, scaleHeight)

        val newWidth = (width * scale).toInt()
        val newHeight = (height * scale).toInt()

        return Bitmap.createScaledBitmap(bitmap, newWidth, newHeight, true)
    }


    private val galleryLauncher = registerForActivityResult(ActivityResultContracts.GetContent()) {
        val uri = it
        if (it != null) {
            val bitmap = try {
                loadBitmapFromUri(it)
            } catch (e: Exception) {
                e.printStackTrace()
                null
            }
            if(bitmap != null) {
                val bef = System.currentTimeMillis()
                lifecycleScope.launch{
                    // Add debugging
                    Log.d("Original", "Bitmap size: ${bitmap.width}x${bitmap.height}")

                    val caption = imageCaptioner.generateCaption(bitmap)
                    val time = System.currentTimeMillis() - bef
                    Log.d("time", time.toString())
                    Toast.makeText(context, caption+time.toString(), Toast.LENGTH_LONG).show()
                    Log.d("Generated Caption", caption)
                }
            }
            Log.d(TAG, uri.toString())
        }
        else{
            Log.i(TAG, "failed to fetch uri")
        }
    }

}

@Composable
fun CameraPreview( cameraController: LifecycleCameraController) {

    // Key Point: Displaying the Camera Preview
    AndroidView(
        modifier = Modifier.fillMaxSize(),
        factory = { ctx ->
            // Initialize the PreviewView and configure it
            PreviewView(ctx).apply {
                scaleType = PreviewView.ScaleType.FILL_START
                implementationMode = PreviewView.ImplementationMode.COMPATIBLE
                controller = cameraController // Set the controller to manage the camera lifecycle
            }
        },
        onRelease = {
            // Release the camera controller when the composable is removed from the screen
            cameraController.unbind()
        }
    )
}
I want to process a frame only after previous one processed, but