namespace DbaMessageInserter;

/// <summary>
/// Canonical Huffman encoder table. Codes are reversed for DEFLATE's
/// least-significant-bit-first bit packing.
/// </summary>
internal sealed class HuffmanTable
{
    private HuffmanTable(uint[] codes, byte[] lengths)
    {
        Codes = codes;
        Lengths = lengths;
    }

    public uint[] Codes { get; }
    public byte[] Lengths { get; }

    public static HuffmanTable FromCodeLengths(
        IReadOnlyList<byte> lengths,
        int maximumAllowedBits)
    {
        if (lengths is null) throw new ArgumentNullException(nameof(lengths));
        if (lengths.Count == 0)
            throw new ArgumentException("A Huffman table must contain at least one symbol.", nameof(lengths));
        if (maximumAllowedBits <= 0 || maximumAllowedBits > 15)
            throw new ArgumentOutOfRangeException(nameof(maximumAllowedBits));

        var copiedLengths = new byte[lengths.Count];
        var bitLengthCounts = new int[maximumAllowedBits + 1];
        int nonZeroCodeCount = 0;

        for (int symbol = 0; symbol < lengths.Count; symbol++)
        {
            byte length = lengths[symbol];
            if (length > maximumAllowedBits)
            {
                throw new ArgumentException(
                    $"Symbol {symbol} has code length {length}, exceeding {maximumAllowedBits}.",
                    nameof(lengths));
            }

            copiedLengths[symbol] = length;
            if (length != 0)
            {
                bitLengthCounts[length]++;
                nonZeroCodeCount++;
            }
        }

        if (nonZeroCodeCount == 0)
            throw new ArgumentException("A Huffman table needs a non-zero code length.", nameof(lengths));

        int left = 1;
        for (int bits = 1; bits <= maximumAllowedBits; bits++)
        {
            left = (left << 1) - bitLengthCounts[bits];
            if (left < 0)
                throw new ArgumentException("The Huffman lengths are oversubscribed.", nameof(lengths));
        }

        var nextCode = new int[maximumAllowedBits + 1];
        int code = 0;
        for (int bits = 1; bits <= maximumAllowedBits; bits++)
        {
            code = (code + bitLengthCounts[bits - 1]) << 1;
            nextCode[bits] = code;
        }

        var reversedCodes = new uint[lengths.Count];
        for (int symbol = 0; symbol < copiedLengths.Length; symbol++)
        {
            int length = copiedLengths[symbol];
            if (length == 0) continue;
            int canonicalCode = nextCode[length]++;
            reversedCodes[symbol] = ReverseLowBits((uint)canonicalCode, length);
        }

        return new HuffmanTable(reversedCodes, copiedLengths);
    }

    public void WriteSymbol(BitWriter writer, int symbol)
    {
        if (writer is null) throw new ArgumentNullException(nameof(writer));
        if (symbol < 0 || symbol >= Lengths.Length)
            throw new ArgumentOutOfRangeException(nameof(symbol));

        int length = Lengths[symbol];
        if (length == 0)
            throw new InvalidOperationException($"Huffman symbol {symbol} has no code.");

        writer.WriteBits(Codes[symbol], length);
    }

    private static uint ReverseLowBits(uint value, int bitCount)
    {
        uint reversed = 0;
        for (int i = 0; i < bitCount; i++)
        {
            reversed = (reversed << 1) | (value & 1u);
            value >>= 1;
        }
        return reversed;
    }
}
