//
//  Deflate.swift
//  SWCompression
//
//  Created by Timofey Solomko on 23.10.16.
//  Copyright © 2017 Timofey Solomko. All rights reserved.
//

import Foundation

/**
 Error happened during deflate decompression. 
 It may indicate that either the data is damaged or it might not be compressed with DEFLATE at all.
 
 - `WrongUncompressedBlockLengths`: `length` and `nlength` bytes of uncompressed block were not compatible.
 - `WrongBlockType`: unsupported block type (not 0, 1 or 2).
 - `WrongSymbol`: unsupported Huffman tree's symbol.
 - `SymbolNotFound`: symbol from input data was not found in Huffman tree.
 */
public enum DeflateError: Error {
    /// Uncompressed block' `length` and `nlength` bytes were not compatible.
    case WrongUncompressedBlockLengths
    /// Unknown block type (not from 0 to 2).
    case WrongBlockType
    /// Decoded symbol was found in Huffman tree but is unknown.
    case WrongSymbol
    /// Symbol was not found in Huffman tree.
    case SymbolNotFound
}

/// Provides function to decompress data, which were compressed with DEFLATE.
public final class Deflate: DecompressionAlgorithm {

    /**
        Decompresses `compressedData` with DEFLATE algortihm.

        If data passed is not actually compressed with DEFLATE, `DeflateError` will be thrown.

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

     - Returns: Decompressed data.
     */
    public static func decompress(compressedData data: Data) throws -> Data {
        /// Object with input data which supports convenient work with bit shifts.
        var pointerData = DataWithPointer(data: data, bitOrder: .reversed)
        return Data(bytes: try decompress(&pointerData))
    }

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

        while true {
            /// Is this a last block?
            let isLastBit = pointerData.bit()
            /// Type of the current block.
            let blockType = [UInt8](pointerData.bits(count: 2).reversed())

            if blockType == [0, 0] { // Uncompressed block.
                pointerData.skipUntilNextByte()
                /// Length of the uncompressed data.
                let length = pointerData.intFromBits(count: 16)
                /// 1-complement of the length.
                let nlength = pointerData.intFromBits(count: 16)
                // Check if lengths are OK (nlength should be a 1-complement of length).
                guard length & nlength == 0 else { throw DeflateError.WrongUncompressedBlockLengths }
                // Process uncompressed data into the output
                for _ in 0..<length {
                    out.append(pointerData.alignedByte())
                }
            } else if blockType == [1, 0] || blockType == [0, 1] {
                // 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: HuffmanTree
                /// Huffman tree for backward distance symbols/codes.
                var mainDistances: HuffmanTree

                if blockType == [0, 1] { // Static Huffman
                    // In this case codes for literals and distances are fixed.
                    // Bootstraps for trees (first element in pair is code, second is number of bits).
                    let staticHuffmanBootstrap = [[0, 8], [144, 9], [256, 7], [280, 8], [288, -1]]
                    let staticHuffmanLengthsBootstrap = [[0, 5], [32, -1]]
                    // Initialize trees from these bootstraps.
                    mainLiterals = HuffmanTree(bootstrap: staticHuffmanBootstrap, &pointerData)
                    mainDistances = HuffmanTree(bootstrap: staticHuffmanLengthsBootstrap, &pointerData)
                } 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).

                    /// Number of literals codes.
                    let literals = pointerData.intFromBits(count: 5) + 257
                    /// Number of distances codes.
                    let distances = pointerData.intFromBits(count: 5) + 1
                    /// Number of code lengths codes.
                    let codeLengthsLength = pointerData.intFromBits(count: 4) + 4

                    // Read code lengths codes.
                    // Moreover, they are stored in a very specific order (defined by HuffmanTree.Constants.codeLengthOrders).
                    var lengthsForOrder = Array(repeating: 0, count: 19)
                    for i in 0..<codeLengthsLength {
                        lengthsForOrder[HuffmanTree.Constants.codeLengthOrders[i]] = pointerData.intFromBits(count: 3)
                    }
                    /// Huffman tree for code lengths. Each code in the main alphabets is coded with this tree.
                    let dynamicCodes = HuffmanTree(lengthsToOrder: lengthsForOrder, &pointerData)

                    // Now we need to read codes (code lengths) for two main alphabets (trees).
                    var codeLengths: [Int] = []
                    var n = 0
                    while n < (literals + distances) {
                        // Finding next Huffman tree's symbol in data.
                        let symbol = dynamicCodes.findNextSymbol()
                        guard symbol != -1 else { throw DeflateError.SymbolNotFound }

                        let count: Int
                        let what: Int
                        if symbol >= 0 && symbol <= 15 {
                            // It is a raw code length.
                            count = 1
                            what = symbol
                        } else if symbol == 16 {
                            // Copy previous code length 3 to 6 times.
                            // Next two bits show how many times we need to copy.
                            count = pointerData.intFromBits(count: 2) + 3
                            what = codeLengths.last!
                        } else if symbol == 17 {
                            // Repeat code length 0 for from 3 to 10 times.
                            // Next three bits show how many times we need to copy.
                            count = pointerData.intFromBits(count: 3) + 3
                            what = 0
                        } else if symbol == 18 {
                            // Repeat code length 0 for from 11 to 138 times.
                            // Next seven bits show how many times we need to do this.
                            count = pointerData.intFromBits(count: 7) + 11
                            what = 0
                        } else {
                            throw DeflateError.WrongSymbol
                        }
                        for _ in 0..<count {
                            codeLengths.append(what)
                        }
                        n += count
                    }
                    // We have read codeLengths for both trees at once.
                    // Now we need to split them and make corresponding trees.
                    mainLiterals = HuffmanTree(lengthsToOrder: Array(codeLengths[0..<literals]), &pointerData)
                    mainDistances = HuffmanTree(lengthsToOrder: Array(codeLengths[literals..<codeLengths.count]), &pointerData)
                }

                // 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.
                        let length = HuffmanTree.Constants.lengthBase[nextSymbol - 257] +
                            pointerData.intFromBits(count: 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.
                        let distance = HuffmanTree.Constants.distanceBase[distanceCode] +
                            pointerData.intFromBits(count: 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 out
    }

    /**
     Compresses `data` with DEFLATE algortihm.

     If during compression something goes wrong `DeflateError` will be thrown.

     - Note: Currently, SWCompression creates only one block for all data
     and the block can either be uncompressed or compressed with static Huffman encoding.
     Uncompressed block is created if amount of data provided is less than 3 bytes and
     static Huffman is used in all other cases.
     */
    public static func compress(data: Data) throws -> Data {
        let bytes = data.toArray(type: UInt8.self)

        if bytes.count < 3 {
            return Data(bytes: Deflate.createUncompressedBlock(bytes))
        }

        let bldCodes = Deflate.lengthEncode(bytes)

        // Let's count possible sizes according to statistics.

        // Uncompressed block size calculation is simple:
        let uncompBlockSize = 1 + 2 + 2 + bytes.count // Header, length, n-length and data.

        // Static Huffman size is more complicated...
        var bitsCount = 3 // Three bits for block's header.
        for (symbol, symbolCount) in bldCodes.stats.enumerated() {
            let codeSize: Int
            // There are extra bits for some codes.
            let extraBitsCount: Int
            switch symbol {
            case 0...143:
                codeSize = 8
                extraBitsCount = 0
            case 144...255:
                codeSize = 9
                extraBitsCount = 0
            case 256...279:
                codeSize = 7
                extraBitsCount = 256 <= symbol && symbol <= 260 ? 0 : (((symbol - 257) >> 2) - 1)
            case 280...285:
                codeSize = 8
                extraBitsCount = symbol == 285 ? 0 : (((symbol - 257) >> 2) - 1)
            case 286...315:
                codeSize = 5
                extraBitsCount = symbol == 286 || symbol == 287 ? 0 : (((symbol - 286) >> 1) - 1)
            default:
                throw DeflateError.SymbolNotFound
            }
            bitsCount += (symbolCount * (codeSize + extraBitsCount))
        }
        let staticHuffmanBlockSize = bitsCount % 8 == 0 ? bitsCount / 8 : bitsCount / 8 + 1

        // Since `length` of uncompressed block is 16-bit integer,
        // there is a limitation on size of uncompressed block.
        // Falling back to static Huffman encoding in case of big uncompressed block is a band-aid solution.
        // TODO: Implement spliting uncompressed block into smaller blocks.
        if uncompBlockSize <= staticHuffmanBlockSize && uncompBlockSize <= 65535 {
            // In case if according to our calculations static huffman will only make output data then input,
            // we fallback to creating uncompressed block.
            // In this case dynamic Huffman encoding can be efficient.
            // TODO: Implement dynamic Huffman code!
            return Data(bytes: Deflate.createUncompressedBlock(bytes))
        } else {
            return Data(bytes: try Deflate.encodeHuffmanBlock(bldCodes.codes))
        }
    }

    private static func createUncompressedBlock(_ bytes: [UInt8]) -> [UInt8] {
        let bitWriter = BitToByteWriter(bitOrder: .reversed)

        // Write block header.
        // Note: Only one block is supported for now.
        bitWriter.write(bit: 1)
        bitWriter.write(bits: [0, 0])

        // Before writing lengths we need to discard remaining bits in current byte.
        bitWriter.finish()

        // Write data's length.
        bitWriter.write(number: bytes.count, bitsCount: 16)
        // Write data's n-length.
        bitWriter.write(number: bytes.count ^ (1 << 16 - 1), bitsCount: 16)

        // Finishing block header.
        bitWriter.finish()
        var out = bitWriter.buffer

        // Write actual data.
        for byte in bytes {
            out.append(byte)
        }

        return out
    }

    private static func encodeHuffmanBlock(_ bldCodes: [BLDCode]) throws -> [UInt8] {
        let bitWriter = BitToByteWriter(bitOrder: .reversed)

        // Write block header.
        // Note: For now it is only static huffman blocks.
        // Note: Only one block is supported for now.
        bitWriter.write(bit: 1)
        bitWriter.write(bits: [1, 0])

        /// Empty DWP object for creating Huffman trees.
        var pointerData = DataWithPointer(data: Data(), bitOrder: .reversed)

        // Constructing Huffman trees for the case of block with preset alphabets.
        // In this case codes for literals and distances are fixed.
        // Bootstraps for trees (first element in pair is code, second is number of bits).
        let staticHuffmanBootstrap = [[0, 8], [144, 9], [256, 7], [280, 8], [288, -1]]
        let staticHuffmanLengthsBootstrap = [[0, 5], [32, -1]]
        /// Huffman tree for literal and length symbols/codes.
        let mainLiterals = HuffmanTree(bootstrap: staticHuffmanBootstrap, &pointerData)
        /// Huffman tree for backward distance symbols/codes.
        let mainDistances = HuffmanTree(bootstrap: staticHuffmanLengthsBootstrap, &pointerData)

        for code in bldCodes {
            switch code {
            case .byte(let byte):
                let codeOfByte = mainLiterals.code(symbol: byte.toInt())
                guard codeOfByte.count > 0
                    else { throw DeflateError.SymbolNotFound }
                bitWriter.write(bits: codeOfByte)
            case .lengthDistance(let length, let distance, let lengthSymbol, let distanceSymbol):
                let lengthExtra = length - HuffmanTree.Constants.lengthBase[lengthSymbol - 257]
                let extraLengthBitsCount = (257 <= lengthSymbol && lengthSymbol <= 260) || lengthSymbol == 285 ?
                    0 : (((lengthSymbol - 257) >> 2) - 1)

                let codeOfLength = mainLiterals.code(symbol: lengthSymbol)
                guard codeOfLength.count > 0
                    else { throw DeflateError.SymbolNotFound }
                bitWriter.write(bits: codeOfLength)
                bitWriter.write(number: lengthExtra, bitsCount: extraLengthBitsCount)

                let distanceExtra = distance - HuffmanTree.Constants.distanceBase[distanceSymbol]
                let extraDistanceBitsCount = distanceSymbol == 0 || distanceSymbol == 1 ? 0 : ((distanceSymbol >> 1) - 1)

                let codeOfDistance = mainDistances.code(symbol: distanceSymbol)
                guard codeOfDistance.count > 0
                    else { throw DeflateError.SymbolNotFound }
                bitWriter.write(bits: codeOfDistance)
                bitWriter.write(number: distanceExtra, bitsCount: extraDistanceBitsCount)
            }
        }

        // End data symbol.
        bitWriter.write(bits: mainLiterals.code(symbol: 256))
        bitWriter.finish()

        return bitWriter.buffer
    }

    private enum BLDCode: CustomStringConvertible {
        case byte(UInt8)
        case lengthDistance(Int, Int, Int, Int)

        var description: String {
            switch self {
            case .byte(let byte):
                return "raw symbol: \(byte)"
            case .lengthDistance(let length, let distance, let lengthSymbol, let distanceSymbol):
                return "length: \(length), length symbol: \(lengthSymbol), distance: \(distance), distance symbol: \(distanceSymbol)"
            }
        }
    }

    private static func lengthEncode(_ rawBytes: [UInt8]) -> (codes: [BLDCode], stats: [Int]) {
        precondition(rawBytes.count >= 3, "Too small array!")

        var buffer: [BLDCode] = []
        var inputIndex = 0
        /// Keys --- three-byte crc32, values --- positions in `rawBytes`.
        var dictionary = [UInt32 : Int]()

        var stats = Array(repeating: 0, count: 316)

        while inputIndex < rawBytes.count {
            let byte = rawBytes[inputIndex]

            // For last two bytes there certainly will be no match.
            // Moreover, `threeByteCrc` cannot be computed, so we need to put them in as `.byte`s.
            // To simplify code we check for this case explicitly.
            if inputIndex >= rawBytes.count - 2 {
                buffer.append(BLDCode.byte(byte))
                stats[byte.toInt()] += 1
                if inputIndex != rawBytes.count - 1 { // For the case of two remaining bytes.
                    buffer.append(BLDCode.byte(rawBytes[inputIndex + 1]))
                    stats[rawBytes[inputIndex + 1].toInt()] += 1
                }
                break
            }

            let threeByteCrc = CheckSums.crc32([rawBytes[inputIndex],
                                                rawBytes[inputIndex + 1],
                                                rawBytes[inputIndex + 2]])

            if let matchStartIndex = dictionary[threeByteCrc] {
                // We need to update position of this match to keep distances as small as possible.
                dictionary[threeByteCrc] = inputIndex

                /// - Note: Minimum match length equals to three.
                var matchLength = 3
                /// Cyclic index which is used to compare bytes in match and in input.
                var repeatIndex = matchStartIndex + matchLength

                /// - Note: Maximum allowed distance equals to 32768.
                let distance = inputIndex - matchStartIndex

                // Again, the distance cannot be greater than 32768.
                if distance <= 32768 {
                    while inputIndex + matchLength < rawBytes.count &&
                        rawBytes[inputIndex + matchLength] == rawBytes[repeatIndex] && matchLength < 258 {
                        matchLength += 1
                        repeatIndex += 1
                        if repeatIndex > inputIndex {
                            repeatIndex = matchStartIndex + 1
                        }
                    }
                    let lengthSymbol = HuffmanTree.Constants.lengthCode[matchLength - 3]
                    let distanceSymbol = ((HuffmanTree.Constants.distanceBase.index { $0 > distance }) ?? 30) - 1
                    buffer.append(BLDCode.lengthDistance(matchLength, distance, lengthSymbol, distanceSymbol))
                    stats[lengthSymbol] += 1
                    stats[286 + distanceSymbol] += 1
                    inputIndex += matchLength
                } else {
                    buffer.append(BLDCode.byte(byte))
                    stats[byte.toInt()] += 1
                    inputIndex += 1
                }
            } else {
                // We need to remember where we met this three-byte sequence.
                dictionary[threeByteCrc] = inputIndex

                buffer.append(BLDCode.byte(byte))
                stats[byte.toInt()] += 1
                inputIndex += 1
            }
            // TODO: Add limitation for dictionary size.
        }

        stats[256] += 1

        return (buffer, stats)
    }

}