package com.example.livingai.data.camera

import android.content.Context
import android.graphics.Bitmap
import android.graphics.RectF
import android.util.Log
import com.example.livingai.R
import com.example.livingai.domain.camera.*
import com.example.livingai.domain.model.camera.*
import com.example.livingai.utils.SignedMask
import com.example.livingai.utils.SilhouetteManager
import com.google.mlkit.vision.common.InputImage
import com.google.mlkit.vision.segmentation.subject.SubjectSegmentation
import com.google.mlkit.vision.segmentation.subject.SubjectSegmenterOptions
import kotlinx.coroutines.suspendCancellableCoroutine
import org.tensorflow.lite.Interpreter
import org.tensorflow.lite.support.common.FileUtil
import java.io.IOException
import java.nio.ByteBuffer
import java.nio.ByteOrder
import java.util.concurrent.atomic.AtomicBoolean
import kotlin.coroutines.resume
import kotlin.math.abs
import kotlin.math.max
import kotlin.math.min

/* ============================================================= */
/* ORIENTATION CHECKER                                           */
/* ============================================================= */

class DefaultOrientationChecker : OrientationChecker {
    override suspend fun analyze(input: PipelineInput): Instruction {

        val isPortraitRequired =
            input.orientation.lowercase() == "front" ||
                    input.orientation.lowercase() == "back"

        val isPortrait = input.deviceOrientation == 90 || input.deviceOrientation == 270
        val isLandscape = input.deviceOrientation == 0 || input.deviceOrientation == 180

        val valid = if (isPortraitRequired) isPortrait else isLandscape

        return Instruction(
            message = if (valid) "Orientation Correct"
            else if (isPortraitRequired) "Turn to portrait mode"
            else "Turn to landscape mode",
            animationResId = if (valid) null else R.drawable.ic_launcher_foreground,
            isValid = valid,
            result = OrientationResult(
                input.deviceOrientation,
                if (isPortraitRequired) CameraOrientation.PORTRAIT else CameraOrientation.LANDSCAPE
            )
        )
    }
}

/* ============================================================= */
/* TILT CHECKER                                                  */
/* ============================================================= */

class DefaultTiltChecker : TiltChecker {
    override suspend fun analyze(input: PipelineInput): Instruction {

        val tolerance = 25f
        var level_message = "Keep the phone straight"

        Log.d("TiltChecker", "Device Roll: ${input.deviceRoll}, Device Pitch: ${input.devicePitch}, Device Azimuth: ${input.deviceAzimuth}")


        val isLevel = when (input.requiredOrientation) {
            CameraOrientation.PORTRAIT ->
                abs(input.devicePitch + 90f) <= tolerance
            CameraOrientation.LANDSCAPE ->
                abs(input.devicePitch) <= tolerance
        }

        val standardPitch = when (input.requiredOrientation) {
            CameraOrientation.PORTRAIT -> -90f
            CameraOrientation.LANDSCAPE -> 0f
        }

        if (!isLevel) {
            if (input.devicePitch > standardPitch) {
                level_message = "Rotate the phone Right"
            } else if (input.devicePitch < standardPitch) {
                level_message = "Rotate the phone Left"
            }
        }

        return Instruction(
            message = if (isLevel) "Device is level" else level_message,
            isValid = isLevel,
            result = TiltResult(input.deviceRoll, input.devicePitch, isLevel)
        )
    }
}

/* ============================================================= */
/* TFLITE OBJECT DETECTOR (PRIMARY + REFERENCE OBJECTS)           */
/* ============================================================= */

class TFLiteObjectDetector(context: Context) : ObjectDetector {

    private var interpreter: Interpreter? = null
    private var labels: List<String> = emptyList()
    private var inputW = 0
    private var inputH = 0
    private var maxDetections = 25

    init {
        try {
            interpreter = Interpreter(
                FileUtil.loadMappedFile(context, "efficientdet-lite0.tflite")
            )
            labels = FileUtil.loadLabels(context, "labels.txt")

            val inputShape = interpreter!!.getInputTensor(0).shape()
            inputW = inputShape[1]
            inputH = inputShape[2]

            maxDetections = interpreter!!.getOutputTensor(0).shape()[1]

        } catch (e: IOException) {
            Log.e("Detector", "Failed to load model", e)
            interpreter = null
        }
    }

    override suspend fun analyze(input: PipelineInput): Instruction {

        val image = input.image
            ?: return Instruction("Waiting for camera", isValid = false)

        val resized = Bitmap.createScaledBitmap(image, inputW, inputH, true)
        val buffer = bitmapToBuffer(resized)

        val locations = Array(1) { Array(maxDetections) { FloatArray(4) } }
        val classes = Array(1) { FloatArray(maxDetections) }
        val scores = Array(1) { FloatArray(maxDetections) }
        val count = FloatArray(1)

        interpreter?.runForMultipleInputsOutputs(
            arrayOf(buffer),
            mapOf(0 to locations, 1 to classes, 2 to scores, 3 to count)
        )

        val detections = mutableListOf<Detection>()

        for (i in 0 until count[0].toInt()) {
            if (scores[0][i] < 0.5f) continue

            val label = labels.getOrElse(classes[0][i].toInt()) { "Unknown" }
            val b = locations[0][i]

            detections += Detection(
                label,
                scores[0][i],
                RectF(
                    b[1] * image.width,
                    b[0] * image.height,
                    b[3] * image.width,
                    b[2] * image.height
                )
            )
        }

        val primary = detections
            .filter { it.label.equals(input.targetAnimal, true) }
            .maxByOrNull { it.confidence }

        val refs = detections
            .filter { it !== primary }
            .mapIndexed { i, d ->
                ReferenceObject(
                    id = "ref_$i",
                    label = d.label,
                    bounds = d.bounds,
                    relativeHeight = d.bounds.height() / image.height,
                    relativeWidth = d.bounds.width() / image.width,
                    distance = null
                )
            }

        return Instruction(
            message = if (primary != null) "Cow detected" else "Cow not detected",
            isValid = primary != null,
            result = DetectionResult(
                isAnimalDetected = primary != null,
                animalBounds = primary?.bounds,
                referenceObjects = refs,
                label = primary?.label,
                confidence = primary?.confidence ?: 0f
            )
        )
    }

    private fun bitmapToBuffer(bitmap: Bitmap): ByteBuffer {
        val buffer = ByteBuffer.allocateDirect(inputW * inputH * 3)
        buffer.order(ByteOrder.nativeOrder())
        val pixels = IntArray(inputW * inputH)
        bitmap.getPixels(pixels, 0, inputW, 0, 0, inputW, inputH)
        for (p in pixels) {
            buffer.put(((p shr 16) and 0xFF).toByte())
            buffer.put(((p shr 8) and 0xFF).toByte())
            buffer.put((p and 0xFF).toByte())
        }
        return buffer
    }

    data class Detection(val label: String, val confidence: Float, val bounds: RectF)
}

class MockPoseAnalyzer : PoseAnalyzer {

    private val segmenter by lazy {
        SubjectSegmentation.getClient(
            SubjectSegmenterOptions.Builder()
                .enableForegroundConfidenceMask()
                .build()
        )
    }

    private val isSegmentationRunning = AtomicBoolean(false)

    private var lastSegmentationValid: Boolean? = null

    override suspend fun analyze(input: PipelineInput): Instruction {

        val detection = input.previousDetectionResult
            ?: return invalidate("No detection")

        val cowBoxImage = detection.animalBounds
            ?: return invalidate("Cow not detected")

        val image = input.image
            ?: return invalidate("No image")

        val silhouette = SilhouetteManager.getSilhouette(input.orientation)
            ?: return invalidate("Silhouette missing")

        val cowBoxScreen = imageToScreenRect(
            box = cowBoxImage,
            imageWidth = image.width,
            imageHeight = image.height,
            screenWidth = input.screenWidthPx,
            screenHeight = input.screenHeightPx
        )

        val silhouetteBoxScreen = silhouette.boundingBox

        val align = checkAlignment(
            detected = cowBoxScreen,
            reference = silhouetteBoxScreen,
            toleranceRatio = 0.15f
        )

        if (align.issue != AlignmentIssue.OK) {
            lastSegmentationValid = false
            return alignmentToInstruction(align)
        }

        // If segmentation already running → reuse last result
        if (!isSegmentationRunning.compareAndSet(false, true)) {

            return Instruction(
                message = if (lastSegmentationValid == true)
                    "Pose Correct"
                else
                    "Hold steady",
                isValid = lastSegmentationValid == true,
                result = detection
            )
        }

        try {
            val cropped = Bitmap.createBitmap(
                image,
                cowBoxImage.left.toInt(),
                cowBoxImage.top.toInt(),
                cowBoxImage.width().toInt(),
                cowBoxImage.height().toInt()
            )

            val resized = Bitmap.createScaledBitmap(
                cropped,
                silhouette.croppedBitmap.width,
                silhouette.croppedBitmap.height,
                true
            )

            val mask = segment(resized)

            val valid = if (mask != null) {
                val score = similarity(mask, silhouette.signedMask)
                score >= 0.40f
            } else {
                false
            }

            lastSegmentationValid = valid

            return Instruction(
                message = if (valid) "Pose Correct" else "Adjust Position",
                isValid = valid,
                result = detection
            )

        } finally {
            isSegmentationRunning.set(false)
        }
    }

    /* -------------------------------------------------- */
    /* HELPERS                                            */
    /* -------------------------------------------------- */

    private fun invalidate(reason: String): Instruction {
        lastSegmentationValid = false
        return Instruction(reason, isValid = false)
    }

    private suspend fun segment(bitmap: Bitmap): ByteArray? =
        suspendCancellableCoroutine { cont ->
            segmenter.process(InputImage.fromBitmap(bitmap, 0))
                .addOnSuccessListener { r ->
                    val buf = r.foregroundConfidenceMask
                        ?: return@addOnSuccessListener cont.resume(null)

                    buf.rewind()
                    val out = ByteArray(bitmap.width * bitmap.height)
                    for (i in out.indices) {
                        out[i] = if (buf.get() > 0.5f) 1 else 0
                    }
                    cont.resume(out)
                }
                .addOnFailureListener {
                    cont.resume(null)
                }
        }

    private fun similarity(mask: ByteArray, ref: SignedMask): Float {
        var s = 0f
        var i = 0
        for (row in ref.mask)
            for (v in row)
                s += mask[i++] * v
        return if (ref.maxValue == 0f) 0f else s / ref.maxValue
    }
}



/* ============================================================= */
/* ALIGNMENT HELPERS (UNCHANGED)                                 */
/* ============================================================= */

enum class AlignmentIssue { TOO_SMALL, TOO_LARGE, MOVE_LEFT, MOVE_RIGHT, MOVE_UP, MOVE_DOWN, OK }

data class AlignmentResult(val issue: AlignmentIssue, val scale: Float, val dx: Float, val dy: Float)

fun checkAlignment(
    detected: RectF,
    reference: RectF,
    toleranceRatio: Float
): AlignmentResult {

    val tolX = reference.width() * toleranceRatio
    val tolY = reference.height() * toleranceRatio

    if (detected.left < reference.left - tolX)
        return AlignmentResult(AlignmentIssue.MOVE_RIGHT, detected.width() / reference.width(), detected.left - reference.left, 0f)

    if (detected.right > reference.right + tolX)
        return AlignmentResult(AlignmentIssue.MOVE_LEFT, detected.width() / reference.width(), detected.right - reference.right, 0f)

    if (detected.top < reference.top - tolY)
        return AlignmentResult(AlignmentIssue.MOVE_DOWN, detected.height() / reference.height(), 0f, detected.top - reference.top)

    if (detected.bottom > reference.bottom + tolY)
        return AlignmentResult(AlignmentIssue.MOVE_UP, detected.height() / reference.height(), 0f, detected.bottom - reference.bottom)

    val scale = min(
        detected.width() / reference.width(),
        detected.height() / reference.height()
    )

    if (scale < 1f - toleranceRatio)
        return AlignmentResult(AlignmentIssue.TOO_SMALL, scale, 0f, 0f)

    if (scale > 1f + toleranceRatio)
        return AlignmentResult(AlignmentIssue.TOO_LARGE, scale, 0f, 0f)

    return AlignmentResult(AlignmentIssue.OK, scale, 0f, 0f)
}

fun alignmentToInstruction(a: AlignmentResult) = when (a.issue) {
    AlignmentIssue.TOO_SMALL -> Instruction("Move closer", false)
    AlignmentIssue.TOO_LARGE -> Instruction("Move backward", false)
    AlignmentIssue.MOVE_LEFT -> Instruction("Move right", false)
    AlignmentIssue.MOVE_RIGHT -> Instruction("Move left", false)
    AlignmentIssue.MOVE_UP -> Instruction("Move down", false)
    AlignmentIssue.MOVE_DOWN -> Instruction("Move up", false)
    AlignmentIssue.OK -> Instruction("Hold steady", true)
}

private fun imageToScreenRect(
    box: RectF,
    imageWidth: Int,
    imageHeight: Int,
    screenWidth: Float,
    screenHeight: Float
): RectF {

    // EXACT SAME LOGIC AS DetectionOverlay
    val widthRatio = screenWidth / imageWidth
    val heightRatio = screenHeight / imageHeight
    val scale = max(widthRatio, heightRatio)

    val offsetX = (screenWidth - imageWidth * scale) / 2f
    val offsetY = (screenHeight - imageHeight * scale) / 2f

    return RectF(
        box.left * scale + offsetX,
        box.top * scale + offsetY,
        box.right * scale + offsetX,
        box.bottom * scale + offsetY
    )
}

/* ============================================================= */
/* CAPTURE + MEASUREMENT (UNCHANGED)                             */
/* ============================================================= */

class DefaultCaptureHandler : CaptureHandler {
    override suspend fun capture(input: PipelineInput, detectionResult: DetectionResult): CaptureData =
        CaptureData(
            image = input.image!!,
            segmentationMask = BooleanArray(0),
            animalMetrics = ObjectMetrics(0f, 0f, 1f),
            referenceObjects = detectionResult.referenceObjects
        )
}

class DefaultMeasurementCalculator : MeasurementCalculator {
    override fun calculateRealMetrics(
        targetHeight: Float,
        referenceObject: ReferenceObject,
        currentMetrics: ObjectMetrics
    ): RealWorldMetrics {

        if (referenceObject.relativeHeight == 0f)
            return RealWorldMetrics(0f, 0f, 0f)

        val scale = targetHeight / referenceObject.relativeHeight

        return RealWorldMetrics(
            height = currentMetrics.relativeHeight * scale,
            width = currentMetrics.relativeWidth * scale,
            distance = currentMetrics.distance
        )
    }
}
