// Copyright (c) 2026 Timofey Solomko
// Licensed under MIT License
//
// See LICENSE for license information

import Foundation
import BitByteData

/// Provides functions for compression and decompression for Deflate algorithm.
public class Deflate: DecompressionAlgorithm {

    /**
     Decompresses `data` using Deflate algortihm.

     - Note: This function is specification compliant.

     - Parameter data: Data compressed with Deflate.

     - Throws: `DeflateError` if unexpected byte (bit) sequence was encountered in `data`.
     It may indicate that either data is damaged or it might not be compressed with Deflate at all.

     - Returns: Decompressed data.
     */
    public static func decompress(data: Data) throws -> Data {
        /// Object with input data which supports convenient work with bit shifts.
        let bitReader = LsbBitReader(data: data)
        return try decompress(bitReader)
    }

    static func decompress(_ bitReader: LsbBitReader) throws -> Data {
        /// An array for storing output data
        var out: [UInt8] = []

        // The smallest possible deflate block consists of 10 bits: `isLastBit`, `blockType` (2 bits), and the
        // end-of-block Huffman-encoded symbol (7 bits).
        guard bitReader.bitsLeft >= 10
            else { throw DeflateError.wrongBlockType }

        while true {
            /// Is this a last block?
            let isLastBit = bitReader.bit()
            /// Type of the current block.
            let blockType = bitReader.int(fromBits: 2)

            if blockType == 0 { // Uncompressed block.
                bitReader.align()
                // The uncompressed block consists at the very least of 32 bits or 4 bytes since they are byte-aligned.
                guard bitReader.bytesLeft >= 4
                    else { throw DeflateError.wrongUncompressedBlockLengths }

                /// Length of the uncompressed data.
                let length = bitReader.uint16()
                /// 1-complement of the length.
                let nlength = bitReader.uint16()
                // Check if lengths are OK (nlength should be a 1-complement of length).
                guard length & nlength == 0
                    else { throw DeflateError.wrongUncompressedBlockLengths }

                guard bitReader.bytesLeft >= length
                    else { throw DeflateError.wrongUncompressedBlockLengths }

                // Process uncompressed data into the output
                for _ in 0..<length {
                    out.append(bitReader.byte())
                }
            } else if blockType == 1 || blockType == 2 {
                // Block with Huffman coding (either static or dynamic)

                // Declaration of Huffman trees which will be populated and used later.
                // There are two alphabets in use and each one needs a Huffman tree.

                /// Huffman tree for literal and length symbols/codes.
                var mainLiterals: DecodingTree
                /// Huffman tree for backward distance symbols/codes.
                var mainDistances: DecodingTree

                if blockType == 1 { // Static Huffman
                    // In this case codes for literals and distances are fixed.
                    // Initialize trees from bootstraps.
                    mainLiterals = DecodingTree(Constants.staticHuffmanLiteralCodes, bitReader)
                    mainDistances = DecodingTree(Constants.staticHuffmanDistanceCodes, bitReader)
                } else { // Dynamic Huffman
                    // In this case there are Huffman codes for two alphabets in data right after block header.
                    // Each code defined by a sequence of code lengths (which are compressed themselves with Huffman).

                    guard bitReader.bitsLeft >= 14
                        else { throw DeflateError.symbolNotFound }

                    /// Number of literals codes.
                    let literals = bitReader.int(fromBits: 5) + 257
                    // According to Deflate specification maximum amount of literal/length codes is 286. However, it
                    // is possible to encode a number larger than this with 5 bits. Specification seems to be silent
                    // about this scenario, so we treat it as a corrupted input.
                    guard literals <= 286
                        else { throw DeflateError.wrongSymbol }
                    /// Number of distances codes.
                    let distances = bitReader.int(fromBits: 5) + 1
                    /// Number of code lengths codes.
                    let codeLengthsCount = bitReader.int(fromBits: 4) + 4

                    guard bitReader.bitsLeft >= 3 * codeLengthsCount
                        else { throw DeflateError.symbolNotFound }

                    // Code lengths of the literal/length and distance Huffman trees are encoded with another Huffman
                    // tree using a special alphabet. The code lengths for this latter tree are presented in a specific
                    // order which is stored in `Constants.codeLengthOrders`: `codeLengthOrders[i]` is the code lengths
                    // alphabet symbol at the i-th place in the sequence.
                    var orderedCodeLengths = Array(repeating: 0, count: 19)
                    for i in 0..<codeLengthsCount {
                        orderedCodeLengths[Constants.codeLengthOrders[i]] = bitReader.int(fromBits: 3)
                    }
                    let dynamicCodes = Code.huffmanCodes(from: Deflate.lengths(from: orderedCodeLengths))
                    /// Huffman tree for code lengths. Each code in the main alphabets is coded with this tree.
                    let dynamicCodeTree = DecodingTree(dynamicCodes, bitReader)

                    // Now we need to read codes (code lengths) for two main alphabets (trees).
                    var codeLengths = Array(repeating: 0, count: literals + distances)
                    var n = codeLengths.startIndex
                    while n < codeLengths.endIndex {
                        // Finding next Huffman tree's symbol in data.
                        let symbol = dynamicCodeTree.findNextSymbol()
                        guard symbol != -1 else { throw DeflateError.symbolNotFound }

                        if symbol >= 0 && symbol <= 15 {
                            // It is a raw code length.
                            codeLengths[n] = symbol
                            n += 1
                        } else if symbol == 16 && n > codeLengths.startIndex {
                            // Copy previous code length 3 to 6 times. Next two bits show how many times we need to copy.
                            // This symbol cannot be the first symbol encoding code lengths, since there is nothing to
                            // copy at this point.
                            guard bitReader.bitsLeft >= 2
                                else { throw DeflateError.symbolNotFound }
                            let copyCount = bitReader.int(fromBits: 2) + 3
                            guard n + copyCount <= codeLengths.endIndex
                                else { throw DeflateError.wrongSymbol }
                            for i in 0..<copyCount {
                                codeLengths[n + i] = codeLengths[n - 1]
                            }
                            n += copyCount
                        } else if symbol == 17 {
                            // Repeat code length 0 from 3 to 10 times. Next 3 bits show how many times we need to copy.
                            // Since `codeLengths` array was already preinitialized with zeros, we just have to skip
                            // ahead.
                            guard bitReader.bitsLeft >= 3
                                else { throw DeflateError.symbolNotFound }
                            n += bitReader.int(fromBits: 3) + 3
                        } else if symbol == 18 {
                            // Repeat code length 0 from 11 to 138 times. Next 7 bits show how many times we need to
                            // copy. Since `codeLengths` array was already preinitialized with zeros, we just have to
                            // skip ahead.
                            guard bitReader.bitsLeft >= 7
                                else { throw DeflateError.symbolNotFound }
                            n += bitReader.int(fromBits: 7) + 11
                        } else {
                            throw DeflateError.wrongSymbol
                        }
                    }
                    // Self-consistency check. Failure here can happen if at some point a zero code length was repeated
                    // too many times indicating a corrupted input.
                    guard n == codeLengths.endIndex
                        else { throw DeflateError.wrongSymbol }
                    // We have read `codeLengths` for both trees at once, so we split them and make corresponding trees.
                    let literalCodes = Code.huffmanCodes(from: Deflate.lengths(from: Array(codeLengths.prefix(upTo: literals))))
                    mainLiterals = DecodingTree(literalCodes, bitReader)
                    let distanceCodes = Code.huffmanCodes(from: Deflate.lengths(from: Array(codeLengths.suffix(from: literals))))
                    mainDistances = DecodingTree(distanceCodes, bitReader)
                }

                // Main loop of data decompression.
                while true {
                    // Read next symbol from data.
                    // It will be either literal symbol or a length of (previous) data we will need to copy.
                    let nextSymbol = mainLiterals.findNextSymbol()
                    guard nextSymbol != -1
                        else { throw DeflateError.symbolNotFound }

                    if nextSymbol >= 0 && nextSymbol <= 255 {
                        // It is a literal symbol so we add it straight to the output data.
                        out.append(nextSymbol.toUInt8())
                    } else if nextSymbol == 256 {
                        // It is a symbol indicating the end of data.
                        break
                    } else if nextSymbol >= 257 && nextSymbol <= 285 {
                        // It is a length symbol.
                        // Depending on the value of nextSymbol there might be additional bits in data,
                        // which we need to add to nextSymbol to get the full length.
                        let extraLength = (257 <= nextSymbol && nextSymbol <= 260) || nextSymbol == 285 ?
                            0 : (((nextSymbol - 257) >> 2) - 1)
                        // Actually, nextSymbol is not a starting value of length,
                        // but an index for special array of starting values.
                        guard bitReader.bitsLeft >= extraLength
                            else { throw DeflateError.symbolNotFound }

                        let length = Constants.lengthBase[nextSymbol - 257] + bitReader.int(fromBits: extraLength)

                        // Then we need to get distance code.
                        let distanceCode = mainDistances.findNextSymbol()
                        guard distanceCode != -1
                            else { throw DeflateError.symbolNotFound }
                        guard distanceCode >= 0 && distanceCode <= 29
                            else { throw DeflateError.wrongSymbol }

                        // Again, depending on the distanceCode's value there might be additional bits in data,
                        // which we need to combine with distanceCode to get the actual distance.
                        let extraDistance = distanceCode == 0 || distanceCode == 1 ? 0 : ((distanceCode >> 1) - 1)
                        // And yes, distanceCode is not a first part of distance but rather an index for special array.
                        guard bitReader.bitsLeft >= extraDistance
                            else { throw DeflateError.symbolNotFound }

                        let distance = Constants.distanceBase[distanceCode] + bitReader.int(fromBits: extraDistance)

                        // We should repeat last 'distance' amount of data.
                        // The amount of times we do this is round(length / distance).
                        // length actually indicates the amount of data we get from this nextSymbol.
                        let repeatCount: Int = length / distance
                        let count = out.count
                        for _ in 0..<repeatCount {
                            for i in count - distance..<count {
                                out.append(out[i])
                            }
                        }
                        // Now we deal with the remainings.
                        if length - distance * repeatCount == distance {
                            for i in out.count - distance..<out.count {
                                out.append(out[i])
                            }
                        } else {
                            for i in out.count - distance..<out.count + length - distance * (repeatCount + 1) {
                                out.append(out[i])
                            }
                        }
                    } else {
                        throw DeflateError.wrongSymbol
                    }
                }

            } else {
                throw DeflateError.wrongBlockType
            }

            // End the cycle if it was the last block.
            if isLastBit == 1 {
                break
            }
        }

        return Data(out)
    }

    private static func lengths(from orderedCodeLengths: [Int]) -> [CodeLength] {
        return orderedCodeLengths.enumerated().map({ CodeLength(symbol: $0, codeLength: $1) })
    }

}