package link.roadio.riff.userstate

import java.nio.ByteBuffer
import java.nio.ByteOrder

internal data class WireField(
    val index: Int,
    val type: String,
    val required: Boolean = false,
)

internal data class WireShape(
    val size: Int,
    val fields: List<WireField>,
)

/** Verify offsets, required fields, UTF-8 and aggregate allocation budgets before model allocation.
 * Authority: https://wiki.271.pics/Riff/1-1-0/formats#semantic-gates
 */
object UserStateWireVerifier {
    fun verify(
        bytes: ByteArray,
        root: String,
    ): Int = verify(ByteBuffer.wrap(bytes), root)

    fun verify(
        bytes: ByteBuffer,
        root: String,
    ): Int =
        Reader(bytes.slice().asReadOnlyBuffer()).run {
            require(buffer.limit() in 8..(64 * 1024 * 1024))
            used = 8
            value(root, relative(0))
            used
        }

    private class Reader(
        bytes: ByteBuffer,
    ) {
        val buffer: ByteBuffer = bytes.order(ByteOrder.LITTLE_ENDIAN)

        // AI-INTENT: Bound traversal by encoded storage, not an accidental queue-size ceiling.
        // Each vector/table may visit fields without allocating its full object graph.
        // https://wiki.271.pics/Riff/LATEST/engineering-verification
        var budget = buffer.limit().toLong() * 4
        var used = 0

        fun span(
            start: Long,
            size: Long,
        ): Int {
            require(start >= 0 && size >= 0 && start <= buffer.limit().toLong() - size) { "Invalid FlatBuffers extent" }
            used = maxOf(used, (start + size).toInt())
            return start.toInt()
        }

        fun short(position: Int) = buffer.getShort(span(position.toLong(), 2)).toInt() and 65535

        fun int(position: Int) = buffer.getInt(span(position.toLong(), 4))

        fun relative(position: Int): Int {
            val offset = int(position).toLong() and 0xffffffffL
            require(offset > 0)
            return span(position.toLong() + offset, 4)
        }

        fun width(type: String): Int =
            when {
                type == "bool" || type == "ubyte" || type.startsWith("enum:") -> {
                    1
                }

                type == "float" || type == "uint" || type == "string" || type.startsWith("[") -> {
                    4
                }

                type in setOf("long", "ulong", "double") -> {
                    8
                }

                else -> {
                    UserStateWireTypes.shapes
                        .getValue(type)
                        .size
                        .takeIf { it > 0 } ?: 4
                }
            }

        fun value(
            type: String,
            position: Int,
            maxItems: Int = 1_000_000,
        ) {
            require(--budget >= 0) { "FlatBuffers traversal budget exceeded" }
            when {
                type == "string" -> {
                    val count = int(position)
                    require(count >= 0)
                    val start = span(position.toLong() + 4, count.toLong() + 1)
                    require(buffer.get(start + count) == 0.toByte())
                    try {
                        val text =
                            buffer.duplicate().apply {
                                position(start)
                                limit(start + count)
                            }
                        Charsets.UTF_8.newDecoder().decode(text)
                    } catch (
                        error: java.nio.charset.CharacterCodingException,
                    ) {
                        throw IllegalArgumentException("Invalid FlatBuffers UTF-8", error)
                    }
                }

                type.startsWith("[") -> {
                    val element = type.substring(1, type.length - 1)
                    val count = int(position)
                    require(count in 0..maxItems && count <= budget)
                    val width = width(element)
                    val start = span(position.toLong() + 4, count.toLong() * width)
                    repeat(count) { index -> field(element, start + index * width) }
                }

                type == "bool" -> {
                    span(position.toLong(), 1)
                    require(buffer.get(position).toInt() in 0..1)
                }

                type.startsWith("enum:") -> {
                    span(position.toLong(), 1)
                    require((buffer.get(position).toInt() and 255) in type.substringAfter(':').split(',').map(String::toInt))
                }

                type in setOf("long", "ulong", "double", "float", "uint", "ubyte") -> {
                    span(position.toLong(), width(type).toLong())
                }

                else -> {
                    val shape = UserStateWireTypes.shapes.getValue(type)
                    if (shape.size > 0) {
                        span(position.toLong(), shape.size.toLong())
                        shape.fields.forEach { value(it.type, position + it.index) }
                    } else {
                        val vtable = span(position.toLong() - int(position), 4)
                        val length = short(vtable)
                        val objectSize = short(vtable + 2)
                        require(length >= 4 && length % 2 == 0 && objectSize >= 4)
                        span(vtable.toLong(), length.toLong())
                        span(position.toLong(), objectSize.toLong())
                        shape.fields.forEach { item ->
                            val entry = 4 + item.index * 2
                            val offset = if (entry < length) short(vtable + entry) else 0
                            if (offset == 0) {
                                require(!item.required) { "Missing required FlatBuffers field" }
                            } else {
                                require(offset >= 4 && offset.toLong() + width(item.type) <= objectSize)
                                val bound =
                                    when {
                                        type == "UserStateCommand" && item.index == 3 -> 4
                                        type == "ListeningHistory" && item.index == 1 -> UserStateValidation.CAPACITY
                                        else -> 1_000_000
                                    }
                                field(item.type, position + offset, bound)
                            }
                        }
                    }
                }
            }
        }

        private fun field(
            type: String,
            position: Int,
            maxItems: Int = 1_000_000,
        ) {
            val indirect = type == "string" || type.startsWith("[") || (UserStateWireTypes.shapes[type]?.size == 0)
            value(type, if (indirect) relative(position) else position, maxItems)
        }
    }
}