package jashkasoft.ithink

import java.nio.ByteBuffer

enum class FrameType (val controlFrame: Boolean, val opcode: Int) {
    /**
     * Regular application level text frame
     */
    TEXT(false, 1),

    /**
     * Regular application level binary frame
     */
    BINARY(false, 2),

    /**
     * Low level close frame type
     */
    CLOSE(true, 8),

    /**
     * Low level ping frame type
     */
    PING(true, 9),

    /**
     * Low level pong frame type
     */
    PONG(true, 0xa);

    companion object {
        private val maxOpcode = FrameType.values().maxBy { it.opcode }!!.opcode

        private val byOpcodeArray = Array(maxOpcode + 1) { op -> FrameType.values().singleOrNull { it.opcode == op } }

        /**
         * Find [FrameType] instance by numeric [opcode]
         * @return a [FrameType] instance or `null` of the [opcode] value is not valid
         */
        operator fun get(opcode: Int): FrameType? = if (opcode in 0..maxOpcode) byOpcodeArray[opcode] else null
    }
}


sealed class Frame(val fin: Boolean, val frameType: FrameType, val buffer: ByteBuffer) {
    private val initialSize = buffer.remaining()

    /**
     * Represents an application level binary frame.
     * In a RAW web socket session a big text frame could be fragmented
     * (separated into several text frames so they have [fin] = false except the last one).
     * Note that usually there is no need to handle fragments unless you have a RAW web socket session.
     */
    class Binary(fin: Boolean, buffer: ByteBuffer) : Frame(fin, FrameType.BINARY, buffer)

    override fun toString() = "Frame $frameType (fin=$fin, buffer len = $initialSize)"
}


fun ByteBuffer.copy(size: Int = remaining()): ByteBuffer {
    return ByteBuffer.allocate(size).apply {
        this@copy.slice().moveTo(this@apply)
        clear()
    }
}

fun ByteBuffer.moveTo(destination: ByteBuffer, limit: Int = Int.MAX_VALUE): Int {
    val size = minOf(limit, remaining(), destination.remaining())
    if (size == remaining()) {
        destination.put(this)
    } else {
        val l = limit()
        limit(position() + size)
        destination.put(this)
        limit(l)
    }
    return size
}

fun Boolean.flagAt(at: Int) = if (this) 1 shl at else 0

infix fun Byte.xor(other: Byte) = toInt().xor(other.toInt()).toByte()

fun ByteBuffer.xor(other: ByteBuffer) {
    val bb = slice()
    val mask = other.slice()
    val maskSize = mask.remaining()

    for (i in 0 until bb.remaining()) {
        bb.put(i, bb.get(i) xor mask[i % maskSize])
    }
}

class Serializer {

    private var frameBody: ByteBuffer? = null
    private var maskBuffer: ByteBuffer? = null

    var masking: Boolean = false

    fun serialize(frame: Frame, buffer: ByteBuffer) {
        while (writeCurrentPayload(buffer)) {
            val mask = masking
            setMaskBuffer(mask)

            val headerSize = estimateFrameHeaderSize(frame, mask)
            if (buffer.remaining() < headerSize) {
                break
            }

            serializeHeader(frame, buffer, mask)
            frameBody = frame.buffer.maskedIfNeeded()

            frameBody!!.array().iterator().forEach { print("$it ") }
            println()
        }
    }

    private fun serializeHeader(frame: Frame, buffer: ByteBuffer, mask: Boolean) {
        val size = frame.buffer.remaining()
        val length1 = when {
            size < 126 -> size
            size <= 0xffff -> 126
            else -> 127
        }

        buffer.put(
                (frame.fin.flagAt(7) or frame.frameType.opcode).toByte()
        )
        buffer.put(
                (mask.flagAt(7) or length1).toByte()
        )

        if (length1 == 126) {
            buffer.putShort(frame.buffer.remaining().toShort())
        } else if (length1 == 127) {
            buffer.putLong(frame.buffer.remaining().toLong())
        }

        maskBuffer?.duplicate()?.moveTo(buffer)
    }

    private fun estimateFrameHeaderSize(f: Frame, mask: Boolean): Int {
        val size = f.buffer.remaining()
        return when {
            size < 126 -> 2
            size <= Short.MAX_VALUE -> 2 + 2
            else -> 2 + 8
        } + maskSize(mask)
    }


    private fun writeCurrentPayload(buffer: ByteBuffer): Boolean {
        val frame = frameBody ?: return true
        frame.moveTo(buffer)
        if (!frame.hasRemaining()) {
            frameBody = null
            return true
        }

        return false
    }

    private fun maskSize(mask: Boolean) = if (mask) 4 else 0

    private fun ByteBuffer.maskedIfNeeded() = maskBuffer?.let { mask -> copy().apply { xor(mask) } } ?: this

    private fun setMaskBuffer(mask: Boolean) {
        if (mask) {
            maskBuffer = ByteBuffer.allocate(4).apply {
                putInt(2.hashCode())
                clear()
            }
        } else {
            maskBuffer = null
        }
    }
}

fun main(args: Array<String>) {

    Serializer().apply {
        val content = "hello".toByteArray()
        val content1 = "world".toByteArray()

        val sendBuffer = ByteBuffer.allocate(content.size + 2)
        val sendBuffer1 = ByteBuffer.allocate(content1.size + 2)
        serialize(Frame.Binary(false, ByteBuffer.wrap(content)), sendBuffer)
        serialize(Frame.Binary(true, ByteBuffer.wrap(content1)), sendBuffer1)

        sendBuffer.flip()
        sendBuffer1.flip()

        println(sendBuffer.array().iterator().forEach { print("$it ") })
        println(sendBuffer1.array().iterator().forEach { print("$it ") })
    }
}