namespace NpncbDiagnostic;

/// <summary>Canonical Huffman decoder for Fixed and Dynamic 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 InvalidDataException("A Huffman alphabet cannot be empty.");
        if (maximumAllowedBits <= 0 || maximumAllowedBits > 15)
            throw new ArgumentOutOfRangeException(nameof(maximumAllowedBits));

        var counts = new int[maximumAllowedBits + 1];
        int nonZeroCount = 0;
        int maximumCodeLength = 0;

        for (int symbol = 0; symbol < lengths.Count; symbol++)
        {
            byte length = lengths[symbol];
            if (length > maximumAllowedBits)
            {
                throw new InvalidDataException(
                    $"Huffman symbol {symbol} has invalid code length {length}.");
            }

            if (length == 0) continue;
            counts[length]++;
            nonZeroCount++;
            maximumCodeLength = Math.Max(maximumCodeLength, length);
        }

        if (nonZeroCount == 0)
            throw new InvalidDataException("A Huffman alphabet has no usable code.");

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

        var nextCode = new int[maximumAllowedBits + 1];
        int code = 0;
        for (int bits = 1; bits <= maximumAllowedBits; bits++)
        {
            code = (code + counts[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 < lengths.Count; symbol++)
        {
            int length = lengths[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 data.");
    }

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