namespace NpncbDiagnostic;

/// <summary>
/// Parses and decodes one raw DEFLATE stream, records physical block types,
/// identifies zero-length Stored blocks (NPNCBs), and reconnects their DBA bits.
/// </summary>
internal static class DeflateScanner
{
    private static readonly int[] LengthBases =
    {
        3, 4, 5, 6, 7, 8, 9, 10,
        11, 13, 15, 17,
        19, 23, 27, 31,
        35, 43, 51, 59,
        67, 83, 99, 115,
        131, 163, 195, 227,
        258
    };

    private static readonly int[] LengthExtraBits =
    {
        0, 0, 0, 0, 0, 0, 0, 0,
        1, 1, 1, 1,
        2, 2, 2, 2,
        3, 3, 3, 3,
        4, 4, 4, 4,
        5, 5, 5, 5,
        0
    };

    private static readonly int[] DistanceBases =
    {
        1, 2, 3, 4,
        5, 7,
        9, 13,
        17, 25,
        33, 49,
        65, 97,
        129, 193,
        257, 385,
        513, 769,
        1025, 1537,
        2049, 3073,
        4097, 6145,
        8193, 12289,
        16385, 24577
    };

    private static readonly int[] DistanceExtraBits =
    {
        0, 0, 0, 0,
        1, 1,
        2, 2,
        3, 3,
        4, 4,
        5, 5,
        6, 6,
        7, 7,
        8, 8,
        9, 9,
        10, 10,
        11, 11,
        12, 12,
        13, 13
    };

    private static readonly int[] CodeLengthAlphabetOrder =
    {
        16, 17, 18, 0, 8, 7, 9, 6, 10, 5,
        11, 4, 12, 3, 13, 2, 14, 1, 15
    };

    private static readonly HuffmanTable FixedLiteralLengthTable =
        HuffmanTable.FromCodeLengths(BuildFixedLiteralLengthLengths(), 15);

    private static readonly HuffmanTable FixedDistanceTable =
        HuffmanTable.FromCodeLengths(Enumerable.Repeat((byte)5, 32).ToArray(), 15);

    public static DeflateScanResult Scan(
        byte[] rawDeflate,
        long expectedUncompressedSize,
        uint expectedCrc32,
        DiagnosticOptions options)
    {
        if (rawDeflate is null) throw new ArgumentNullException(nameof(rawDeflate));
        if (expectedUncompressedSize < 0)
            throw new ArgumentOutOfRangeException(nameof(expectedUncompressedSize));
        if (options is null) throw new ArgumentNullException(nameof(options));
        options.Validate();

        var reader = new BitReader(rawDeflate);
        var output = new DeflateOutputTracker(expectedUncompressedSize);
        var dbaBits = new List<byte>();
        var npncbBlocks = new List<NpncbBlockFinding>();

        int physicalBlockCount = 0;
        int storedDataBlockCount = 0;
        int fixedBlockCount = 0;
        int dynamicBlockCount = 0;
        int nonZeroDbaBitCount = 0;
        bool isFinal;

        do
        {
            if (physicalBlockCount >= options.MaximumPhysicalBlocksPerEntry)
                throw new InvalidDataException("The DEFLATE stream contains too many physical blocks.");

            long blockStartBit = reader.BitPosition;
            isFinal = reader.ReadBits(1) != 0;
            int btype = checked((int)reader.ReadBits(2));

            switch (btype)
            {
                case 0:
                {
                    bool wasNpncb = ReadStoredBlock(
                        reader,
                        output,
                        physicalBlockCount,
                        isFinal,
                        blockStartBit,
                        dbaBits,
                        npncbBlocks,
                        options.MaximumDbaBitsPerEntry,
                        ref nonZeroDbaBitCount);

                    if (!wasNpncb) storedDataBlockCount++;
                    break;
                }

                case 1:
                    fixedBlockCount++;
                    ReadCompressedBlock(reader, output, FixedLiteralLengthTable, FixedDistanceTable);
                    break;

                case 2:
                    dynamicBlockCount++;
                    (HuffmanTable literalLength, HuffmanTable distance) = ReadDynamicTables(reader);
                    ReadCompressedBlock(reader, output, literalLength, distance);
                    break;

                default:
                    throw new InvalidDataException("DEFLATE BTYPE=11 is reserved and invalid.");
            }

            physicalBlockCount++;
        }
        while (!isFinal);

        int trailingBitCount = checked((int)reader.RemainingBits);
        if (trailingBitCount > 7)
        {
            throw new InvalidDataException(
                $"The raw DEFLATE entry has {trailingBitCount} trailing bits after the final block.");
        }

        int nonZeroTrailingBits = 0;
        for (int i = 0; i < trailingBitCount; i++)
            nonZeroTrailingBits += reader.ReadBits(1) == 0 ? 0 : 1;

        output.VerifyCompleted(expectedCrc32);

        return new DeflateScanResult(
            physicalBlockCount,
            storedDataBlockCount,
            fixedBlockCount,
            dynamicBlockCount,
            npncbBlocks.Count,
            dbaBits.Count,
            nonZeroDbaBitCount,
            dbaBits,
            npncbBlocks,
            output.Length,
            output.Crc32,
            trailingBitCount,
            nonZeroTrailingBits);
    }

    private static bool ReadStoredBlock(
        BitReader reader,
        DeflateOutputTracker output,
        int physicalBlockIndex,
        bool isFinal,
        long blockStartBit,
        ICollection<byte> allDbaBits,
        ICollection<NpncbBlockFinding> npncbBlocks,
        int maximumDbaBits,
        ref int totalNonZeroDbaBits)
    {
        long dbaStartBit = reader.BitPosition;
        int dbaBitCount = reader.BitsUntilByteBoundary;
        byte dbaValue = 0;
        int nonZeroDbaBits = 0;

        for (int bitIndex = 0; bitIndex < dbaBitCount; bitIndex++)
        {
            byte bit = checked((byte)reader.ReadBits(1));
            dbaValue |= checked((byte)(bit << bitIndex));
            nonZeroDbaBits += bit;
        }

        ushort length = reader.ReadUInt16LittleEndian();
        ushort complement = reader.ReadUInt16LittleEndian();
        if ((ushort)(length ^ 0xFFFF) != complement)
            throw new InvalidDataException("Stored block LEN/NLEN mismatch.");

        bool isNpncb = length == 0;
        if (isNpncb)
        {
            if (allDbaBits.Count > maximumDbaBits - dbaBitCount)
                throw new InvalidDataException("The connected NPNCB DBA exceeds the configured limit.");

            for (int bitIndex = 0; bitIndex < dbaBitCount; bitIndex++)
                allDbaBits.Add((byte)((dbaValue >> bitIndex) & 1));

            totalNonZeroDbaBits = checked(totalNonZeroDbaBits + nonZeroDbaBits);
            npncbBlocks.Add(new NpncbBlockFinding(
                physicalBlockIndex,
                isFinal,
                blockStartBit,
                dbaStartBit,
                dbaBitCount,
                dbaValue,
                nonZeroDbaBits));
        }

        for (int i = 0; i < length; i++)
            output.WriteStoredByte(reader.ReadAlignedByte());

        return isNpncb;
    }

    private static void ReadCompressedBlock(
        BitReader reader,
        DeflateOutputTracker output,
        HuffmanTable literalLengthTable,
        HuffmanTable distanceTable)
    {
        while (true)
        {
            int symbol = literalLengthTable.ReadSymbol(reader);

            if (symbol < 256)
            {
                output.WriteLiteral(checked((byte)symbol));
                continue;
            }

            if (symbol == 256) return;
            if (symbol < 257 || symbol > 285)
                throw new InvalidDataException($"Invalid literal/length symbol {symbol}.");

            int lengthIndex = symbol - 257;
            int length = checked(
                LengthBases[lengthIndex] +
                (int)reader.ReadBits(LengthExtraBits[lengthIndex]));

            int distanceSymbol = distanceTable.ReadSymbol(reader);
            if (distanceSymbol < 0 || distanceSymbol > 29)
                throw new InvalidDataException($"Invalid distance symbol {distanceSymbol}.");

            int distance = checked(
                DistanceBases[distanceSymbol] +
                (int)reader.ReadBits(DistanceExtraBits[distanceSymbol]));

            output.CopyMatch(distance, length);
        }
    }

    private static (HuffmanTable LiteralLength, HuffmanTable Distance) ReadDynamicTables(
        BitReader reader)
    {
        int literalLengthCount = checked(257 + (int)reader.ReadBits(5));
        int distanceCount = checked(1 + (int)reader.ReadBits(5));
        int codeLengthCount = checked(4 + (int)reader.ReadBits(4));

        var codeLengthLengths = new byte[19];
        for (int i = 0; i < codeLengthCount; i++)
            codeLengthLengths[CodeLengthAlphabetOrder[i]] = checked((byte)reader.ReadBits(3));

        HuffmanTable codeLengthTable = HuffmanTable.FromCodeLengths(codeLengthLengths, 7);
        var combinedLengths = new byte[checked(literalLengthCount + distanceCount)];
        int index = 0;

        while (index < combinedLengths.Length)
        {
            int symbol = codeLengthTable.ReadSymbol(reader);
            if (symbol <= 15)
            {
                combinedLengths[index++] = checked((byte)symbol);
                continue;
            }

            int repeatCount;
            byte repeatedValue;

            switch (symbol)
            {
                case 16:
                    if (index == 0)
                        throw new InvalidDataException("Code-length symbol 16 has no previous value.");
                    repeatCount = checked(3 + (int)reader.ReadBits(2));
                    repeatedValue = combinedLengths[index - 1];
                    break;

                case 17:
                    repeatCount = checked(3 + (int)reader.ReadBits(3));
                    repeatedValue = 0;
                    break;

                case 18:
                    repeatCount = checked(11 + (int)reader.ReadBits(7));
                    repeatedValue = 0;
                    break;

                default:
                    throw new InvalidDataException($"Invalid code-length symbol {symbol}.");
            }

            if (index > combinedLengths.Length - repeatCount)
                throw new InvalidDataException("Repeated code lengths exceed the Dynamic header.");

            for (int i = 0; i < repeatCount; i++)
                combinedLengths[index++] = repeatedValue;
        }

        byte[] literalLengthLengths = combinedLengths[..literalLengthCount];
        byte[] distanceLengths = combinedLengths[literalLengthCount..];

        if (literalLengthLengths[256] == 0)
            throw new InvalidDataException("The Dynamic tree has no end-of-block symbol.");

        return (
            HuffmanTable.FromCodeLengths(literalLengthLengths, 15),
            HuffmanTable.FromCodeLengths(distanceLengths, 15));
    }

    private static byte[] BuildFixedLiteralLengthLengths()
    {
        var lengths = new byte[288];
        for (int symbol = 0; symbol <= 143; symbol++) lengths[symbol] = 8;
        for (int symbol = 144; symbol <= 255; symbol++) lengths[symbol] = 9;
        for (int symbol = 256; symbol <= 279; symbol++) lengths[symbol] = 7;
        for (int symbol = 280; symbol <= 287; symbol++) lengths[symbol] = 8;
        return lengths;
    }
}
