import java.math.BigInteger
import kotlin.math.pow

/**
 * A 9P message with the given type, tag, and fields. The message size is calculated automatically.
 *
 * Important note: the field names in [fieldValuesInt] and [fieldValuesStr] (i.e. the keys of their maps) must be
 * mutually exclusive and the union of these two maps' keys must result in a subset of (or a set equal to)
 * [fieldNames]. Calling [write] when these conditions are not met throws an exception.
 *
 * @param type The 9P message type.
 * @param tag The tag given to the message.
 * @param fieldNames The names of the message fields, in the same order they are expected to be sent.
 * @param fieldValuesInt A map of each field name into its value. This map only stores integer values.
 * @param fieldValuesStr A map of each field name into its value. This map only stores string values.
 */
class NinePMessage(val type: NinePMessageType, val tag: UShort, val fieldNames: List<String>, val fieldValuesInt: Map<String, BigInteger>, val fieldValuesStr: Map<String, String>, val maxSize: UInt) {
    /**
     * Intersection between [fieldNames] and [fieldValuesInt]. In other words: the amount of integer fields that are
     * going to be used when writing the message.
     */
    private val insecInts = fieldNames.intersect(fieldValuesInt.keys)
    /**
     * Intersection between [fieldNames] and [fieldValuesStr]. In other words: the amount of string fields that are
     * going to be used when writing the message.
     */
    private val insecStrs = fieldNames.intersect(fieldValuesStr.keys)

    /**
     * Send the message using the given networking API.
     *
     * @param npt The networking API.
     * @throws IllegalArgumentException if [fieldNames], [fieldValuesInt], and [fieldValuesStr] are incoherent or the
     *         final size of the message exceeds the negotiated value.
     */
    fun write(npt: NetworkPacketTransporter) {
        // check that names in fieldNames exist as keys in either fieldValuesInt or fieldValuesStr but not both
        require(insecInts.size == fieldNames.size - insecStrs.size)

        val totalSize = size()
        if (totalSize > this.maxSize) {
            throw IllegalArgumentException("Message size exceeded.")
        }
        writeMessageSizeType(npt, totalSize, type)
        writeInteger(npt, tag.toInt().toBigInteger())
        for (field in fieldNames) {
            if (field in insecInts) {
                writeInteger(npt, fieldValuesInt[field]!!)
            } else {
                writeString(npt, fieldValuesStr[field]!!)
            }
        }
    }

    /**
     * Write an integer number to the connection.
     *
     * In 9P, binary numbers (non-textual) are specified in little-endian order (least significant byte first).
     *
     * @param npt The networking API.
     * @param value The number's value. [SizedMessageField] defines both its actual value and its size.
     */
    private fun writeInteger(npt: NetworkPacketTransporter, value: BigInteger) {
        val bytes = value.toByteArray()
        npt.transmit(Array(bytes.size) { i -> bytes[i].toUByte() })
    }

    /**
     * Write a string to the connection.
     *
     * In 9P, strings are represented as a 2-byte integer (the string's size) followed by the actual UTF-8 string. The
     * null terminator is forbidden in 9P messages.
     *
     * @param npt The networking API.
     * @param value The string.
     * @throws IllegalArgumentException if the value of the string's size does not fit into 2 bytes.
     */
    private fun writeString(npt: NetworkPacketTransporter, value: String) {
        require(value.length <= 2.0.pow(16.0) - 1)
        writeInteger(npt, value.length.toBigInteger())
        val bytes = value.toByteArray()
        npt.transmit(Array(bytes.size) { i -> bytes[i].toUByte() })
    }

    /**
     * Write the message size and type.
     *
     * @param npt The networking API.
     * @param size The total message size, including the 4 bytes of this parameter and the type's byte.
     * @param type The 9P message type as a [NinePMessageType] constant.
     */
    private fun writeMessageSizeType(npt: NetworkPacketTransporter, size: UInt, type: NinePMessageType) {
        writeInteger(npt, size.toInt().toBigInteger())
        writeInteger(npt, type.value.toInt().toBigInteger())
    }

    /**
     * Calculate the expected size of the message.
     */
    fun size(): UInt {
        return 4u + 1u + 2u + this.insecInts.sumOf { this.fieldValuesInt[it]!!.bitLength().toUInt() } + this.insecStrs.sumOf { 2u + this.fieldValuesStr[it]!!.length.toUInt() }
    }
}