package com.performancecoachlab.posedetection.custom

import androidx.compose.runtime.Composable
import androidx.compose.ui.platform.LocalContext
import co.touchlab.kermit.Logger
import org.json.JSONObject
import android.os.Build
import org.tensorflow.lite.Interpreter
import org.tensorflow.lite.gpu.GpuDelegate
import org.tensorflow.lite.support.common.FileUtil
import org.tensorflow.lite.support.metadata.MetadataExtractor
import java.io.ByteArrayInputStream
import java.nio.MappedByteBuffer
import java.nio.charset.StandardCharsets
import java.util.zip.GZIPInputStream
import java.util.zip.Inflater
import java.util.zip.InflaterInputStream
import java.util.zip.ZipInputStream

@Composable
actual fun initialiseObjectModel(modelPath: ModelPath): ObjectModel {
    if (modelPath.androidModelPath == null) {
        throw IllegalArgumentException("Android model path cannot be null")
    }
    // Prefer GPU, then NNAPI (API 27+), then CPU.
    val (options, gpuDelegate) = runCatching {
        val delegate = GpuDelegate()
        val opts = Interpreter.Options().apply {
            addDelegate(delegate)
            setNumThreads(2)
        }
        Logger.d { "TFLite GPU delegate available" }
        opts to delegate
    }.onFailure { t ->
        Logger.w(t) { "TFLite GPU delegate not available; trying NNAPI" }
    }.getOrElse {
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.O_MR1) {
            runCatching {
                val nnapiDelegate = org.tensorflow.lite.nnapi.NnApiDelegate()
                val opts = Interpreter.Options().apply {
                    addDelegate(nnapiDelegate)
                    setNumThreads(2)
                }
                Logger.d { "TFLite NNAPI delegate available" }
                opts to null
            }.onFailure { t ->
                Logger.w(t) { "TFLite NNAPI delegate not available; falling back to CPU" }
            }.getOrElse {
                Interpreter.Options().apply { setNumThreads(4) } to null
            }
        } else {
            Interpreter.Options().apply { setNumThreads(4) } to null
        }
    }

    val model = FileUtil.loadMappedFile(LocalContext.current, modelPath.androidModelPath)
    val labels = labels(model)

    val interpreter = runCatching {
        Interpreter(model, options)
    }.onFailure { t ->
        // If GPU interpreter creation fails, retry on CPU.
        Logger.w(t) { "TFLite GPU Failed to create GPU TFLite interpreter; retrying on CPU" }
        gpuDelegate?.close()
    }.getOrElse {
        val cpuOptions = Interpreter.Options().apply { setNumThreads(4) }
        Interpreter(model, cpuOptions)
    }

    val inputShape = interpreter.getInputTensor(0)?.shape()
    val outputShape = interpreter.getOutputTensor(0)?.shape()
    val modelInfo = ModelInfo.fromShapes(
        inputShape = inputShape
            ?: throw IllegalArgumentException("Invalid model: input shape is null"),
        outputShape = outputShape
            ?: throw IllegalArgumentException("Invalid model: output shape is null"),
        labels,
    )
    val androidDetector = AndroidDetector(
        interpreter = interpreter, modelInfo = modelInfo
    )
    return ObjectModel(androidDetector)
}

fun labels(model: MappedByteBuffer): List<String> {
    return runCatching {
        val extractor = MetadataExtractor(model)
        val files = extractor.associatedFileNames.orEmpty()
        if (files.isEmpty()) return@runCatching emptyList()

        // Try every associated file. Pick the first that decodes to JSON with a `names` object.
        for (name in files) {
            val rawBytes = runCatching { extractor.getAssociatedFile(name).readBytes() }.getOrNull()
                ?: continue

            val decoded = rawBytes.decodeUtf8PossiblyCompressed()
            val trimmed = decoded.trimStart()
            if (!trimmed.startsWith("{")) continue

            val labels = parseUltralyticsNamesJson(decoded)
            if (labels.isNotEmpty()) {
                Logger.d { "Loaded ${labels.size} labels from associated file '$name'" }
                return@runCatching labels
            }
        }
        emptyList()
    }.onFailure { t ->
        Logger.w(t) { "Failed to load labels from TFLite metadata" }
    }.getOrDefault(emptyList())
}

private fun ByteArray.decodeUtf8PossiblyCompressed(): String {
    if (isEmpty()) return ""

    fun firstBytesHex(n: Int = 8): String =
        take(minOf(size, n)).joinToString(" ") { b -> "%02x".format(b) }

    // 0) Scan for embedded magic headers / common compressed stream signatures.
    val zipOffset =
        indexOfSubsequence(byteArrayOf('P'.code.toByte(), 'K'.code.toByte(), 0x03, 0x04))
    val gzipOffset = indexOfSubsequence(byteArrayOf(0x1F.toByte(), 0x8B.toByte()))

    // Common zlib headers (CMF/FLG). Most common are 0x78 0x9C (default), 0x78 0xDA (best), 0x78 0x01 (no compression).
    val zlibOffsets = listOf(
        indexOfSubsequence(byteArrayOf(0x78.toByte(), 0x9C.toByte())),
        indexOfSubsequence(byteArrayOf(0x78.toByte(), 0xDA.toByte())),
        indexOfSubsequence(byteArrayOf(0x78.toByte(), 0x01.toByte())),
    ).filter { it >= 0 }.distinct().sorted()

    if (zipOffset > 0) {
        runCatching {
            val sliced = copyOfRange(zipOffset, size)
            val text = sliced.decodeFromZipIfPossible(preferEntryName = "metadata.json")
            if (text.isNotBlank() && text.trimStart().startsWith("{")) return text
        }
    }

    if (gzipOffset > 0) {
        runCatching {
            val sliced = copyOfRange(gzipOffset, size)
            val text = sliced.decodeFromGzipOrTarGz()
            if (text.isNotBlank() && text.trimStart().startsWith("{")) return text
        }
    }

    for (off in zlibOffsets) {
        if (off <= 0) continue
        runCatching {
            val sliced = copyOfRange(off, size)
            val text = sliced.decodeFromZlib()
            if (text.isNotBlank() && text.trimStart().startsWith("{")) return text
        }
    }

    runCatching {
        val text = decodeFromGzipOrTarGz()
        if (text.isNotBlank() && text.trimStart().startsWith("{")) return text
    }

    runCatching {
        val text = decodeFromZipIfPossible(preferEntryName = "metadata.json")
        if (text.isNotBlank() && text.trimStart().startsWith("{")) return text
    }

    runCatching {
        val text = decodeFromZlib()
        if (text.isNotBlank() && text.trimStart().startsWith("{")) return text
    }

    runCatching {
        val text = decodeFromDeflateRaw()
        if (text.isNotBlank() && text.trimStart().startsWith("{")) return text
    }

    return toString(StandardCharsets.UTF_8)
}

private fun ByteArray.decodeFromDeflateRaw(): String {
    return runCatching {
        val inflater = Inflater(true) // nowrap=true => raw DEFLATE
        InflaterInputStream(ByteArrayInputStream(this), inflater)
            .bufferedReader(StandardCharsets.UTF_8)
            .use { it.readText() }
    }.getOrDefault("")
}

private fun ByteArray.decodeFromZlib(): String {
    return runCatching {
        InflaterInputStream(ByteArrayInputStream(this))
            .bufferedReader(StandardCharsets.UTF_8)
            .use { it.readText() }
    }.getOrDefault("")
}

private fun ByteArray.decodeFromGzipOrTarGz(): String {
    // First try: plain gzipped UTF-8 directly
    runCatching {
        val text = GZIPInputStream(ByteArrayInputStream(this))
            .bufferedReader(StandardCharsets.UTF_8)
            .use { it.readText() }
        if (text.isNotBlank()) return text
    }

    // Second try: gzipped TAR that contains metadata.json (or first file)
    return runCatching {
        GZIPInputStream(ByteArrayInputStream(this)).use { gz ->
            extractFirstTarFileUtf8(gz)
        }
    }.getOrElse {
        ""
    }
}

private fun extractFirstTarFileUtf8(input: java.io.InputStream): String {
    // Minimal TAR reader: TAR headers are 512 bytes.
    // We only need to extract the first regular file (or metadata.json if present)
    // and decode it as UTF-8.

    fun readExactly(buf: ByteArray): Boolean {
        var off = 0
        while (off < buf.size) {
            val r = input.read(buf, off, buf.size - off)
            if (r <= 0) return false
            off += r
        }
        return true
    }

    val header = ByteArray(512)
    var firstText: String? = null

    while (true) {
        if (!readExactly(header)) break
        // End of archive: two consecutive 512-byte blocks of zero
        if (header.all { it == 0.toByte() }) break

        val name = header.copyOfRange(0, 100).toString(StandardCharsets.US_ASCII)
            .trimEnd { it == '\u0000' }
        val sizeOctal = header.copyOfRange(124, 136).toString(StandardCharsets.US_ASCII).trim()
            .trimEnd { it == '\u0000' }
        val typeFlag = header[156]

        val fileSize = sizeOctal.toLongOrNull(8) ?: 0L
        val isRegularFile = typeFlag == 0.toByte() || typeFlag == '0'.code.toByte()

        val fileData = ByteArray(fileSize.toInt())
        if (fileSize > 0 && !readExactly(fileData)) break

        // TAR pads file data to 512 byte blocks.
        val pad = ((512 - (fileSize % 512)) % 512).toInt()
        if (pad > 0) {
            val skip = ByteArray(pad)
            if (!readExactly(skip)) break
        }

        if (isRegularFile) {
            val text = fileData.toString(StandardCharsets.UTF_8)
            if (name.equals("metadata.json", ignoreCase = true)) return text
            if (firstText == null && text.isNotBlank()) firstText = text
        }
    }

    return firstText.orEmpty()
}

private fun ByteArray.indexOfSubsequence(needle: ByteArray): Int {
    if (needle.isEmpty() || size < needle.size) return -1
    outer@ for (i in 0..(size - needle.size)) {
        for (j in needle.indices) {
            if (this[i + j] != needle[j]) continue@outer
        }
        return i
    }
    return -1
}

private fun parseUltralyticsNamesJson(text: String): List<String> {
    return runCatching {
        val root = JSONObject(text)
        val namesObj = root.optJSONObject("names") ?: return@runCatching emptyList()

        val keys = namesObj.keys().asSequence()
            .mapNotNull { k -> k.toIntOrNull()?.let { idx -> idx to k } }
            .sortedBy { it.first }
            .toList()

        if (keys.isEmpty()) return@runCatching emptyList()

        val maxIdx = keys.maxOf { it.first }
        val out = MutableList(maxIdx + 1) { "" }

        for ((idx, key) in keys) {
            val label = namesObj.optString(key, "").trim()
            if (label.isNotBlank() && idx in out.indices) {
                out[idx] = label
            }
        }

        out.filter { it.isNotBlank() }
    }.getOrDefault(emptyList())
}

private fun parseUltralyticsNames(text: String): List<String> {
    // Grab the `names: { ... }` section (non-greedy) to avoid matching other maps.
    val namesBlock = Regex("""['"]names['"]\s*:\s*\{([\s\S]*?)\}""")
        .find(text)
        ?.groupValues
        ?.getOrNull(1)
        ?: return emptyList()

    // Match entries like: 0: 'person'  OR  0: "person"
    val entry = Regex("""(\d+)\s*:\s*['"]([^'"]+)['"]""")
    val pairs = entry.findAll(namesBlock).mapNotNull { m ->
        val idx = m.groupValues[1].toIntOrNull() ?: return@mapNotNull null
        val name = m.groupValues[2].trim()
        if (name.isBlank()) null else idx to name
    }.toList()

    if (pairs.isEmpty()) return emptyList()

    val maxIdx = pairs.maxOf { it.first }
    val out = MutableList(maxIdx + 1) { "" }
    for ((i, label) in pairs) {
        if (i in out.indices) out[i] = label
    }
    return out.filter { it.isNotBlank() }
}

data class ModelInfo(
    val inputShape: IntArray,
    val outputShape: IntArray,
    val inputWidth: Int,
    val inputHeight: Int,
    val inputChannels: Int? = null,
    val isNhwc: Boolean? = null,
    val labels: List<String>,
) {
    companion object {
        fun fromShapes(
            inputShape: IntArray,
            outputShape: IntArray,
            labels: List<String>
        ): ModelInfo {
            // Common TFLite image shapes:
            // NHWC: [1, H, W, C]
            // NCHW: [1, C, H, W] (less common on Android)
            // Some models might be [H, W, C] (no batch).
            val (h, w, c, nhwc) = when (inputShape.size) {
                4 -> {
                    val isNhwcGuess = inputShape[3] in 1..4
                    if (isNhwcGuess) {
                        Quad(inputShape[1], inputShape[2], inputShape[3], true)
                    } else {
                        Quad(inputShape[2], inputShape[3], inputShape[1], false)
                    }
                }

                3 -> Quad(inputShape[0], inputShape[1], inputShape[2], true)
                else -> Quad(0, 0, null, null)
            }

            return ModelInfo(
                inputShape = inputShape,
                outputShape = outputShape,
                inputWidth = w,
                inputHeight = h,
                inputChannels = c,
                isNhwc = nhwc,
                labels = labels,
            )
        }
    }
}

private data class Quad(
    val h: Int, val w: Int, val c: Int?, val nhwc: Boolean?
)

data class AndroidDetector(
    val interpreter: Interpreter, val modelInfo: ModelInfo
)

actual class ObjectModel {

    private var detector: AndroidDetector? = null

    constructor(detector: AndroidDetector) {
        this.detector = detector
    }

    fun getDetector(): AndroidDetector? {
        return detector
    }
}

private fun ByteArray.decodeFromZipIfPossible(preferEntryName: String): String {
    return runCatching {
        ZipInputStream(ByteArrayInputStream(this)).use { zis ->
            var firstNonEmpty: String? = null
            val entries = mutableListOf<String>()

            while (true) {
                val entry = zis.nextEntry ?: break
                entries += entry.name
                if (entry.isDirectory) continue

                val entryBytes = zis.readBytes()
                val text = entryBytes.toString(StandardCharsets.UTF_8)

                if (entry.name.equals(preferEntryName, ignoreCase = true)) {
                    Logger.d { "TFLite metadata zip entry=${entry.name} size=${entryBytes.size} entries=$entries" }
                    return@use text
                }

                if (firstNonEmpty == null && text.isNotBlank()) {
                    firstNonEmpty = text
                }
            }

            firstNonEmpty.orEmpty()
        }
    }.getOrDefault("")
}

@Composable
internal actual fun platformRememberObjectModel(modelPath: ModelPath): ObjectModel {
    return ObjectModelProvider.get(modelPath)
}