package com.performancecoachlab.posedetection.camera

import android.graphics.Bitmap
import androidx.annotation.OptIn
import androidx.camera.core.CameraSelector.DEFAULT_BACK_CAMERA
import androidx.camera.core.CameraSelector.DEFAULT_FRONT_CAMERA
import androidx.camera.core.ExperimentalGetImage
import androidx.camera.view.PreviewView
import androidx.compose.foundation.Image
import androidx.compose.foundation.layout.Box
import androidx.compose.foundation.layout.fillMaxSize
import androidx.compose.foundation.layout.height
import androidx.compose.foundation.layout.padding
import androidx.compose.foundation.layout.width
import androidx.compose.material3.Surface
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.setValue
import androidx.compose.ui.Modifier
import androidx.compose.ui.draw.scale
import androidx.compose.ui.graphics.ImageBitmap
import androidx.compose.ui.graphics.asAndroidBitmap
import androidx.compose.ui.graphics.asImageBitmap
import androidx.compose.ui.layout.ContentScale
import androidx.compose.ui.platform.LocalContext
import androidx.compose.ui.unit.dp
import androidx.compose.ui.viewinterop.AndroidView
import androidx.lifecycle.LifecycleOwner
import androidx.lifecycle.compose.LocalLifecycleOwner
import com.google.android.gms.tasks.Tasks
import com.google.mlkit.vision.common.InputImage
import com.google.mlkit.vision.pose.Pose
import com.google.mlkit.vision.pose.PoseDetection
import com.google.mlkit.vision.pose.PoseLandmark
import com.google.mlkit.vision.pose.defaults.PoseDetectorOptions
import com.performancecoachlab.posedetection.skeleton.Skeleton
import com.performancecoachlab.posedetection.skeleton.SkeletonRepository
import java.util.concurrent.Executors
import java.io.File
import androidx.camera.lifecycle.ProcessCameraProvider
import androidx.camera.core.Preview
import androidx.camera.core.ImageAnalysis
import androidx.camera.core.ImageProxy
import androidx.compose.runtime.DisposableEffect
import androidx.compose.runtime.rememberCoroutineScope
import androidx.compose.ui.geometry.Rect
import com.google.mlkit.vision.pose.PoseDetector
import com.performancecoachlab.posedetection.custom.CustomObjectRespository
import com.performancecoachlab.posedetection.custom.ObjectModel
import com.performancecoachlab.posedetection.recording.AnalysisObject
import com.performancecoachlab.posedetection.recording.AnalysisResult
import kotlinx.coroutines.launch
import org.tensorflow.lite.support.image.TensorImage
import org.tensorflow.lite.task.vision.detector.ObjectDetector

// Data class to hold recording state for each recording ID
data class RecordingSlot(
    var builder: com.performancecoachlab.posedetection.encoding.VideoBuilder? = null,
    var firstTimestampMs: Long? = null,
    val outputPath: String
)

@OptIn(ExperimentalGetImage::class)
@Composable
actual fun CameraView(
    skeletonRepository: SkeletonRepository,
    customObjectRepository: CustomObjectRespository,
    detectMode: DetectMode,
    drawSkeleton: Boolean,
    objectModel: ObjectModel?,
    drawObjects: ((List<AnalysisObject>) -> List<DrawableObject>)?,
    modifier: Modifier,
    frontCamera: Boolean,
    recordingId: String?,
    focusArea: Rect?,
    controller: CameraViewController?,
    onRecordToggled: (Boolean) -> Unit,
    onVideoSaved: (String, String) -> Unit,
) {
    var bitmap by remember { mutableStateOf<ImageBitmap?>(null) }
    val options =
        PoseDetectorOptions.Builder().setDetectorMode(PoseDetectorOptions.STREAM_MODE).build()
    val poseDetector = PoseDetection.getClient(options)
    val context = LocalContext.current
    val lifecycleOwner: LifecycleOwner = LocalLifecycleOwner.current
    val previewView: PreviewView = remember { PreviewView(context) }
    val executor = remember { Executors.newSingleThreadExecutor() }
    val scope = rememberCoroutineScope()
    var focus by remember { mutableStateOf(focusArea) }
    var objectDetector by remember { mutableStateOf(objectModel) }
    var currentDetectMode by remember { mutableStateOf(detectMode) }
    LaunchedEffect(detectMode) { currentDetectMode = detectMode }

    // Update focus when focusArea changes
    LaunchedEffect(focusArea) {
        focus = focusArea
    }

    // Update object model when it changes
    LaunchedEffect(objectModel) {
        objectDetector = objectModel
    }

    // Multi-recording state - map of recording ID to recording slot
    var activeRecordings by remember { mutableStateOf<Map<String, RecordingSlot>>(emptyMap()) }

    suspend fun startRecording(id: String, width: Int, height: Int) {
        val outputPath = File(context.cacheDir, "camera_recording_${id}_${System.currentTimeMillis()}.mp4").absolutePath
        val slot = RecordingSlot(
            builder = null, // Lazily created on first frame
            firstTimestampMs = null,
            outputPath = outputPath
        )
        activeRecordings = activeRecordings + (id to slot)
    }

    suspend fun stopRecording(id: String) {
        activeRecordings[id]?.let { slot ->
            slot.builder?.let { builder ->
                try {
                    val path = builder.finalize()
                    onVideoSaved(id, path)
                } catch (e: Exception) {
                    //println("Error finalizing recording $id: ${e.message}")
                }
            }
            activeRecordings = activeRecordings - id
        }
    }

    fun addFrameToActiveRecordings(frame: ImageBitmap, timestampMs: Long) {
        if (activeRecordings.isNotEmpty()) {
            scope.launch {
                activeRecordings.forEach { (id, slot) ->
                    try {
                        // Lazily create builder on first frame to lock width/height
                        if (slot.builder == null) {
                            slot.builder = com.performancecoachlab.posedetection.encoding.createVideoBuilder(
                                outputPath = slot.outputPath,
                                fps = 30,
                                width = frame.width,
                                height = frame.height
                            )
                            slot.firstTimestampMs = timestampMs
                        }

                        slot.builder?.let { builder ->
                            val relativeTimestamp = timestampMs - (slot.firstTimestampMs ?: timestampMs)
                            builder.addFrame(frame, relativeTimestamp)
                        }
                    } catch (e: Exception) {
                        // Handle recording errors gracefully
                        //println("Error recording frame for ID $id: ${e.message}")
                    }
                }
            }
        }
    }

    // React to recordingIds changes - diff and start/stop recordings
    LaunchedEffect(recordingId) {
        val currentIds = activeRecordings.keys
        recordingId?.also { id ->
            currentIds.forEach {
                if(!it.contentEquals(id)){
                    stopRecording(it)
                }
            }
            if(!currentIds.contains(id)){
                if (bitmap != null) {
                    startRecording(id, bitmap!!.width, bitmap!!.height)
                } else {
                    // If no frame yet, create empty slot that will be initialized on first frame
                    val outputPath = File(context.cacheDir, "camera_recording_${id}_${System.currentTimeMillis()}.mp4").absolutePath
                    val slot = RecordingSlot(
                        builder = null,
                        firstTimestampMs = null,
                        outputPath = outputPath
                    )
                    activeRecordings = activeRecordings + (id to slot)
                }
            }
        }?: run {
            // No recordingId - stop all
            currentIds.forEach {
                stopRecording(it)
            }
        }


    }

    LaunchedEffect(lifecycleOwner) {
        val cameraProvider = ProcessCameraProvider.getInstance(context).get()
        cameraProvider.unbindAll()
        val cameraSelector = if (frontCamera) DEFAULT_FRONT_CAMERA else DEFAULT_BACK_CAMERA
        val preview = Preview.Builder().build().also {
            it.surfaceProvider = previewView.surfaceProvider
        }
        val imageAnalysis = ImageAnalysis.Builder().build().also { analysis ->
            analysis.setAnalyzer(executor) { imageProxy ->
                val timestamp = System.currentTimeMillis()
                val area = focus
                val poseClient   = if (currentDetectMode.doPose()) poseDetector else null
                val objectClient = if (currentDetectMode.doObject()) objectDetector?.getDetector() else null
                imageProxy.process(
                    objectClient, poseClient, timestamp, area
                ){ analysisResult, _bitmap ->
                    customObjectRepository.updateCustomObject(analysisResult.objects)
                    analysisResult.skeleton?.let { skel ->
                        skeletonRepository.updateSkeleton(skel)
                    }
                    bitmap = _bitmap.asImageBitmap().let { inbmp ->
                                controller?.setRequestDataProvider {
                                    CameraViewData(
                                        width = inbmp.width.toFloat(),
                                        height = inbmp.height.toFloat(),
                                        rotation = when (imageProxy.imageInfo.rotationDegrees) {
                                            0 -> SensorRotation.ROTATION_0
                                            90 -> SensorRotation.ROTATION_90
                                            180 -> SensorRotation.ROTATION_180
                                            270 -> SensorRotation.ROTATION_270
                                            else -> SensorRotation.ROTATION_0
                                        }
                                    )
                                }
                                addFrameToActiveRecordings(inbmp, timestamp)
                                inbmp.drawResults(if(drawSkeleton) analysisResult.skeleton else null, drawObjects?.invoke(analysisResult.objects)?: emptyList())
                            }
                    imageProxy.close()
                }
            }
        }
        cameraProvider.bindToLifecycle(
            lifecycleOwner,
            cameraSelector,
            preview,
            imageAnalysis
        )
        previewView.scaleType = PreviewView.ScaleType.FIT_CENTER
    }
    Box(
        modifier = modifier
    ) {
        Box(
            modifier = Modifier
                .height(100.dp)
                .width(100.dp)
                .padding(20.dp)
        ) {
            AndroidView(
                factory = { previewView }, modifier = Modifier.fillMaxSize()
            )

        }
        Surface(modifier = Modifier.fillMaxSize()) {
            bitmap?.also {
                Image(
                    it,
                    contentDescription = "video feed",
                    modifier = Modifier
                        .fillMaxSize()
                        .scale(if (frontCamera) -1f else 1f, 1f),
                    contentScale = ContentScale.Crop
                )
            }
        }
    }
}



private fun PoseLandmark?.toSkeletonCoords(): Skeleton.SkeletonCoordinate? {
    return this?.position?.let {
        Skeleton.SkeletonCoordinate(
            x = it.x,
            y = it.y,
        )
    }
}

fun Bitmap.rotate(angle: Float): Bitmap {
    val matrix = android.graphics.Matrix()
    matrix.postRotate(angle)
    return Bitmap.createBitmap(this, 0, 0, width, height, matrix, true)
}

fun ImageBitmap.rotate(angle: Float): ImageBitmap {
    return asAndroidBitmap().rotate(angle).asImageBitmap()
}

fun skeleton(pose: Pose?, timestamp: Long, width: Int, height: Int): Skeleton {
    return Skeleton(
        timestamp = timestamp,
        leftShoulder = pose?.getPoseLandmark(PoseLandmark.LEFT_SHOULDER)?.toSkeletonCoords(),
        rightShoulder = pose?.getPoseLandmark(PoseLandmark.RIGHT_SHOULDER)?.toSkeletonCoords(),
        leftElbow = pose?.getPoseLandmark(PoseLandmark.LEFT_ELBOW)?.toSkeletonCoords(),
        rightElbow = pose?.getPoseLandmark(PoseLandmark.RIGHT_ELBOW)?.toSkeletonCoords(),
        leftWrist = pose?.getPoseLandmark(PoseLandmark.LEFT_WRIST)?.toSkeletonCoords(),
        rightWrist = pose?.getPoseLandmark(PoseLandmark.RIGHT_WRIST)?.toSkeletonCoords(),
        leftHip = pose?.getPoseLandmark(PoseLandmark.LEFT_HIP)?.toSkeletonCoords(),
        rightHip = pose?.getPoseLandmark(PoseLandmark.RIGHT_HIP)?.toSkeletonCoords(),
        leftKnee = pose?.getPoseLandmark(PoseLandmark.LEFT_KNEE)?.toSkeletonCoords(),
        rightKnee = pose?.getPoseLandmark(PoseLandmark.RIGHT_KNEE)?.toSkeletonCoords(),
        leftAnkle = pose?.getPoseLandmark(PoseLandmark.LEFT_ANKLE)?.toSkeletonCoords(),
        rightAnkle = pose?.getPoseLandmark(PoseLandmark.RIGHT_ANKLE)?.toSkeletonCoords(),
        width = width.toFloat(),
        height = height.toFloat(),
    )
}
