namespace DbaMessageRestorer;

/// <summary>
/// Canonical Huffman decoder table for walking raw DEFLATE blocks.
/// </summary>
internal sealed class HuffmanTable
{
    private readonly Dictionary<uint, int>[] _decodeByLength;
    private readonly int _maximumCodeLength;

    private HuffmanTable(
        Dictionary<uint, int>[] decodeByLength,
        int maximumCodeLength)
    {
        _decodeByLength = decodeByLength;
        _maximumCodeLength = maximumCodeLength;
    }

    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;
        int maximumCodeLength = 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++;
                maximumCodeLength = Math.Max(maximumCodeLength, length);
            }
        }

        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 decodeByLength = new Dictionary<uint, int>[maximumAllowedBits + 1];
        for (int bits = 0; bits < decodeByLength.Length; bits++)
            decodeByLength[bits] = new Dictionary<uint, int>();

        for (int symbol = 0; symbol < copiedLengths.Length; symbol++)
        {
            int length = copiedLengths[symbol];
            if (length == 0) continue;
            int canonicalCode = nextCode[length]++;
            uint reversedCode = ReverseLowBits((uint)canonicalCode, length);
            decodeByLength[length].Add(reversedCode, symbol);
        }

        return new HuffmanTable(decodeByLength, maximumCodeLength);
    }

    public int ReadSymbol(BitReader reader)
    {
        if (reader is null) throw new ArgumentNullException(nameof(reader));
        uint reversedCode = 0;

        for (int length = 1; length <= _maximumCodeLength; length++)
        {
            reversedCode |= reader.ReadBits(1) << (length - 1);
            if (_decodeByLength[length].TryGetValue(reversedCode, out int symbol))
                return symbol;
        }

        throw new InvalidDataException("Invalid Huffman code in raw DEFLATE.");
    }

    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;
    }
}
