using System.Buffers.Binary;
using System.Text;

namespace DbaMessageRestorer;

public sealed record NpncbDbaReadLog(
    int PhysicalBlockIndex,
    bool IsFinal,
    long BlockStartBit,
    long DbaStartBit,
    int DbaBitCount);

public sealed record DbaZipMessageResult(
    string ZipPath,
    string EntryName,
    int CompressedSize,
    int PhysicalBlockCount,
    int NpncbCount,
    int DbaBitCount,
    IReadOnlyList<byte> DbaBits,
    IReadOnlyList<NpncbDbaReadLog> NpncbBlocks,
    DbaDecodedMessage DecodedMessage);

public static class ZipDbaMessageReader
{
    public static DbaZipMessageResult Read(string zipPath)
    {
        if (string.IsNullOrWhiteSpace(zipPath))
            throw new ArgumentException("ZIP path is required.", nameof(zipPath));

        RawZipEntry entry = RawZipEntryReader.ReadFirstDeflateEntry(zipPath);
        DeflateDbaScanResult scan = DeflateDbaScanner.Scan(entry.RawDeflateBytes);
        DbaDecodedMessage decoded = DbaMessageDecoder.Decode(scan.DbaBits);

        return new DbaZipMessageResult(
            zipPath,
            entry.EntryName,
            entry.RawDeflateBytes.Length,
            scan.PhysicalBlockCount,
            scan.NpncbBlocks.Count,
            scan.DbaBits.Count,
            scan.DbaBits,
            scan.NpncbBlocks,
            decoded);
    }
}

internal sealed record RawZipEntry(string EntryName, byte[] RawDeflateBytes);

internal static class RawZipEntryReader
{
    private const uint LocalHeaderSignature = 0x0403_4B50;
    private const uint CentralHeaderSignature = 0x0201_4B50;
    private const uint EocdSignature = 0x0605_4B50;
    private const ushort DeflateMethod = 8;
    private const ushort Utf8Flag = 1 << 11;
    private const ushort EncryptedFlag = 1;

    public static RawZipEntry ReadFirstDeflateEntry(string zipPath)
    {
        byte[] zip = File.ReadAllBytes(zipPath);
        int eocdOffset = FindEndOfCentralDirectory(zip);
        ushort thisDisk = ReadUInt16(zip, eocdOffset + 4);
        ushort centralStartDisk = ReadUInt16(zip, eocdOffset + 6);
        ushort entriesOnDisk = ReadUInt16(zip, eocdOffset + 8);
        ushort totalEntries = ReadUInt16(zip, eocdOffset + 10);
        uint centralSize = ReadUInt32(zip, eocdOffset + 12);
        uint centralOffsetValue = ReadUInt32(zip, eocdOffset + 16);

        if (thisDisk != 0 || centralStartDisk != 0 || entriesOnDisk != totalEntries)
            throw new NotSupportedException("Multi-disk ZIP is not supported.");
        if (totalEntries == 0) throw new InvalidDataException("ZIP has no entries.");
        if (centralSize == uint.MaxValue || centralOffsetValue == uint.MaxValue)
            throw new NotSupportedException("ZIP64 is not supported.");

        int centralOffset = checked((int)centralOffsetValue);
        EnsureRange(zip, centralOffset, 46);
        if (ReadUInt32(zip, centralOffset) != CentralHeaderSignature)
            throw new InvalidDataException("Invalid central-directory header.");

        ushort flags = ReadUInt16(zip, centralOffset + 8);
        ushort method = ReadUInt16(zip, centralOffset + 10);
        uint compressedSizeValue = ReadUInt32(zip, centralOffset + 20);
        ushort fileNameLength = ReadUInt16(zip, centralOffset + 28);
        ushort extraLength = ReadUInt16(zip, centralOffset + 30);
        ushort commentLength = ReadUInt16(zip, centralOffset + 32);
        ushort diskNumberStart = ReadUInt16(zip, centralOffset + 34);
        uint localOffsetValue = ReadUInt32(zip, centralOffset + 42);

        if ((flags & EncryptedFlag) != 0)
            throw new NotSupportedException("Encrypted entries are not supported.");
        if (method != DeflateMethod)
            throw new NotSupportedException($"Entry method {method} is not DEFLATE method 8.");
        if (compressedSizeValue == uint.MaxValue || localOffsetValue == uint.MaxValue)
            throw new NotSupportedException("ZIP64 entry fields are not supported.");
        if (diskNumberStart != 0)
            throw new NotSupportedException("Entry starts on another disk.");

        int centralVariableLength = checked(fileNameLength + extraLength + commentLength);
        EnsureRange(zip, centralOffset + 46, centralVariableLength);
        ReadOnlySpan<byte> entryNameBytes = zip.AsSpan(centralOffset + 46, fileNameLength);
        string entryName = (flags & Utf8Flag) != 0
            ? Encoding.UTF8.GetString(entryNameBytes)
            : DecodeLatin1(entryNameBytes);

        int localOffset = checked((int)localOffsetValue);
        EnsureRange(zip, localOffset, 30);
        if (ReadUInt32(zip, localOffset) != LocalHeaderSignature)
            throw new InvalidDataException("Invalid local ZIP header.");

        ushort localNameLength = ReadUInt16(zip, localOffset + 26);
        ushort localExtraLength = ReadUInt16(zip, localOffset + 28);
        int dataOffset = checked(localOffset + 30 + localNameLength + localExtraLength);
        int compressedSize = checked((int)compressedSizeValue);
        EnsureRange(zip, dataOffset, compressedSize);

        return new RawZipEntry(entryName, zip.AsSpan(dataOffset, compressedSize).ToArray());
    }

    private static string DecodeLatin1(ReadOnlySpan<byte> bytes)
    {
        var characters = new char[bytes.Length];
        for (int i = 0; i < bytes.Length; i++) characters[i] = (char)bytes[i];
        return new string(characters);
    }

    private static int FindEndOfCentralDirectory(byte[] zip)
    {
        const int minimumEocdLength = 22;
        if (zip.Length < minimumEocdLength)
            throw new InvalidDataException("File is too short to be ZIP.");

        int minimumOffset = Math.Max(0, zip.Length - minimumEocdLength - ushort.MaxValue);
        for (int offset = zip.Length - minimumEocdLength; offset >= minimumOffset; offset--)
        {
            if (ReadUInt32(zip, offset) != EocdSignature) continue;
            ushort commentLength = ReadUInt16(zip, offset + 20);
            if (offset + minimumEocdLength + commentLength == zip.Length) return offset;
        }

        throw new InvalidDataException("EOCD record was not found.");
    }

    private static ushort ReadUInt16(byte[] bytes, int offset)
    {
        EnsureRange(bytes, offset, 2);
        return BinaryPrimitives.ReadUInt16LittleEndian(bytes.AsSpan(offset, 2));
    }

    private static uint ReadUInt32(byte[] bytes, int offset)
    {
        EnsureRange(bytes, offset, 4);
        return BinaryPrimitives.ReadUInt32LittleEndian(bytes.AsSpan(offset, 4));
    }

    private static void EnsureRange(byte[] bytes, int offset, int length)
    {
        if (offset < 0 || length < 0 || offset > bytes.Length - length)
            throw new InvalidDataException("ZIP structure points outside file bounds.");
    }
}

internal sealed record DeflateDbaScanResult(
    int PhysicalBlockCount,
    List<byte> DbaBits,
    List<NpncbDbaReadLog> NpncbBlocks);

internal static class DeflateDbaScanner
{
    private const int MaximumPhysicalBlocks = 2_000_000;

    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[] 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 DeflateDbaScanResult Scan(byte[] rawDeflate)
    {
        if (rawDeflate is null) throw new ArgumentNullException(nameof(rawDeflate));
        var reader = new BitReader(rawDeflate);
        var dbaBits = new List<byte>();
        var npncbBlocks = new List<NpncbDbaReadLog>();
        int physicalBlockCount = 0;
        bool isFinal;

        do
        {
            if (physicalBlockCount >= MaximumPhysicalBlocks)
                throw new InvalidDataException("Too many physical DEFLATE blocks.");

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

            switch (btype)
            {
                case 0:
                    ReadStoredBlock(reader, physicalBlockCount, isFinal, blockStartBit, dbaBits, npncbBlocks);
                    break;
                case 1:
                    SkipCompressedBlock(reader, FixedLiteralLengthTable, FixedDistanceTable);
                    break;
                case 2:
                    (HuffmanTable literalLength, HuffmanTable distance) = ReadDynamicTables(reader);
                    SkipCompressedBlock(reader, literalLength, distance);
                    break;
                default:
                    throw new InvalidDataException("DEFLATE BTYPE=11 is invalid.");
            }

            physicalBlockCount++;
        }
        while (!isFinal);

        return new DeflateDbaScanResult(physicalBlockCount, dbaBits, npncbBlocks);
    }

    private static void ReadStoredBlock(
        BitReader reader,
        int physicalBlockIndex,
        bool isFinal,
        long blockStartBit,
        ICollection<byte> allDbaBits,
        ICollection<NpncbDbaReadLog> npncbBlocks)
    {
        long dbaStartBit = reader.BitPosition;
        int dbaBitCount = reader.BitsUntilByteBoundary;
        var candidateDbaBits = new byte[dbaBitCount];
        for (int i = 0; i < dbaBitCount; i++) candidateDbaBits[i] = (byte)reader.ReadBits(1);

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

        if (length == 0)
        {
            foreach (byte bit in candidateDbaBits) allDbaBits.Add(bit);
            npncbBlocks.Add(new NpncbDbaReadLog(
                physicalBlockIndex, isFinal, blockStartBit, dbaStartBit, dbaBitCount));
        }

        reader.SkipAlignedBytes(length);
    }

    private static void SkipCompressedBlock(
        BitReader reader,
        HuffmanTable literalLengthTable,
        HuffmanTable distanceTable)
    {
        while (true)
        {
            int symbol = literalLengthTable.ReadSymbol(reader);
            if (symbol < 256) continue;
            if (symbol == 256) return;
            if (symbol < 257 || symbol > 285)
                throw new InvalidDataException($"Invalid literal/length symbol {symbol}.");

            reader.ReadBits(LengthExtraBits[symbol - 257]);
            int distanceSymbol = distanceTable.ReadSymbol(reader);
            if (distanceSymbol < 0 || distanceSymbol > 29)
                throw new InvalidDataException($"Invalid distance symbol {distanceSymbol}.");
            reader.ReadBits(DistanceExtraBits[distanceSymbol]);
        }
    }

    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]] = (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++] = (byte)symbol;
                continue;
            }

            int repeatCount;
            byte repeatedValue;
            switch (symbol)
            {
                case 16:
                    if (index == 0) throw new InvalidDataException("Code-length 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 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("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;
    }
}
