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

import Foundation
import Testing
import BitByteData
import SWCompression

struct DeflateTests {

    private static let testType: String = "deflate"

    @Test func randomInputTruncations() throws {
        // In this test we check that there is no crash when dealing with the truncation in the middle of the Deflate
        // compressed data. The idea is to take three different types of Deflate blocks (uncompressed, static Huffman,
        // and dynamic Huffman), truncate the input data manually at a random point inside it, and then test if an
        // appropriate error is thrown. To make test a bit more sophisticated we generate a number of random truncations
        // for each tested file.

        // test6 contains dynamic Huffman Deflate block.
        // test8 contains static Huffman Deflate block.
        // test9 contains uncompressed Deflate block.

        for testName in ["test6", "test8", "test9"] {
            let testData = try Constants.data(forTest: testName, withType: DeflateTests.testType)
            for _ in 0..<10 {
                let truncationIndex = Int.random(in: (testData.startIndex + 1)..<testData.endIndex)
                #expect(throws: (any Error).self, "No error thrown, \(testName), truncationIndex=\(truncationIndex)") {
                    try Deflate.decompress(data: testData[..<truncationIndex])
                }
            }
        }
    }

    @Test func symbol16First() {
        // Symbol 16 cannot be the first symbol encoding code lengths in the dynamic Huffman block. Previously, there
        // was a crash in such situation. This test checks that error is thrown appropriately.

        // The input data was constructed manually:
        // - 101: last block bit and dynamical Huffman block type,
        // - 00000: minimal amount of literal codes,
        // - 00000: minimal amount of distances codes,
        // - 0000: minimal amount of code length codes (4),
        // - 011 011 010 001: four code lenghts for code length codes (1, 2, 3, 3),
        // - 0: encoded code length symbol 16 (which means to copy previous code length).
        let testData = Data([0b0000_0101, 0b0000_0000, 0b1010_0010, 0b0000_1101])
        #expect(throws: (any Error).self) { try Deflate.decompress(data: testData) }
    }

    @Test func codeLengthsOverCopy() {
        // Previously, when decoding code lengths in a dynamic Huffman block if a copy code length count was
        // sufficiently large, it could lead to an inconsistent state or a crash (due to out-of-range array subscript).
        // In the new version these situations are checked.

        // Repeat a zero code length too many times. The input data was constructed manually:
        // - 101: last block bit and dynamical Huffman block type,
        // - 00000: minimal amount of literal codes,
        // - 00000: minimal amount of distances codes,
        // - 0000: minimal amount of code length codes (4),
        // - 011 011 010 001: four code lenghts for code length codes (1, 2, 3, 3),
        // - 111: encoded code length symbol 18 (which means to repeat 0 code length),
        // - 1111111: maximum repeat amount,
        // - 111: encoded code length symbol 18 (which means to repeat 0 code length),
        // - 1111111: maximum repeat amount.
        var testData = Data([0b0000_0101, 0b0000_0000, 0b1010_0010, 0b1110_1101, 0b1111_1111, 0b1111_1111, 0b0000_0001])
        #expect(throws: (any Error).self) { try Deflate.decompress(data: testData) }

        // Copy a previous code length too many times. The input data was constructed manually:
        // - 101: last block bit and dynamical Huffman block type,
        // - 00000: minimal amount of literal codes,
        // - 00000: minimal amount of distances codes,
        // - 0000: minimal amount of code length codes (4),
        // - 011 011 010 001: four code lenghts for code length codes (1, 2, 3, 3),
        // - 111: encoded code length symbol 18 (which means to repeat 0 code length),
        // - 1111111: maximum repeat amount,
        // - 111: encoded code length symbol 18 (which means to repeat 0 code length),
        // - 1101100: repeat amount,
        // - 0: encoded code length symbol 16 (which means to copy a previous code length),
        // - 01: copy amount.
        testData = Data([0b0000_0101, 0b0000_0000, 0b1010_0010, 0b1110_1101, 0b1111_1111, 0b1011_0011, 0b0000_0101])
        #expect(throws: (any Error).self) { try Deflate.decompress(data: testData) }
    }

    @Test(.bug("https://github.com/tsolomko/SWCompression/issues/57", id: 57))
    func codeLengthsAllZero() throws {
        // Previously, if all code lengths for a particular Huffman tree in a dynamic Huffman block were zero, the
        // Huffman code construction algorithm would crash. This behavior was reported as issue #57. After fixing the
        // crash all zero code lengths results in an "empty" Huffman tree. If such a tree is used to read a symbol, no
        // symbol will be found, and an error will be thrown appropriately.

        // Code lengths for the code lengths tree are all zero which indicates corrputed input. The input data was
        // constructed manually:
        // - 101: last block bit and dynamical Huffman block type,
        // - 00000: minimal amount of literal codes,
        // - 00000: minimal amount of distances codes,
        // - 0000: minimal amount of code length codes (4),
        // - 000 000 000 000: four zero code lenghts for code length codes.
        let testData = Data([0b0000_0101, 0b0000_0000, 0b0000_0000, 0b0_0000])
        #expect(throws: (any Error).self) { try Deflate.decompress(data: testData) }

        // The following inputs are based on test6 compressed data, but otherwise were constructed manually with the
        // help of `LsbBitWriter` from BitByteData.

        // All code lengths are zero for the literals/lengths tree, and two non-zero distance tree code lengths.
        var writer = LsbBitWriter()
        // - 101: last block bit and dynamical Huffman block type,
        writer.write(bits: [1, 0, 1])
        // - 11101: maximal amount of literal codes (29 + 257 = 286),
        writer.write(unsignedNumber: 29, bitsCount: 5)
        // - 00001: number of distances codes (1 + 1 = 2),
        writer.write(unsignedNumber: 1, bitsCount: 5)
        // - 1110: number of code length codes (4 + 14 = 18),
        writer.write(unsignedNumber: 14, bitsCount: 4)
        // - 000 011 010 000 000 000 000 000 000 000 000 000 000 010 000 011 000 010: code lenghts for code length codes,
        writer.write(unsignedNumber: 0, bitsCount: 3)
        writer.write(unsignedNumber: 3, bitsCount: 3)
        writer.write(unsignedNumber: 2, bitsCount: 3)
        for _ in 0..<10 {
            writer.write(unsignedNumber: 0, bitsCount: 3)
        }
        writer.write(unsignedNumber: 2, bitsCount: 3)
        writer.write(unsignedNumber: 0, bitsCount: 3)
        writer.write(unsignedNumber: 3, bitsCount: 3)
        writer.write(unsignedNumber: 0, bitsCount: 3)
        writer.write(unsignedNumber: 2, bitsCount: 3)
        // - 01: encoded code length symbol 18 (which means to repeat 0 code length),
        writer.write(unsignedNumber: 1, bitsCount: 2)
        // - 1111111: maximum repeat amount,
        writer.write(unsignedNumber: 127, bitsCount: 7)
        // - 01: encoded code length symbol 18 (which means to repeat 0 code length),
        writer.write(unsignedNumber: 1, bitsCount: 2)
        // - 1111111: maximum repeat amount,
        writer.write(unsignedNumber: 127, bitsCount: 7)
        // - 111: encoded code length symbol 17 (which means to repeat 0 code length),
        writer.write(unsignedNumber: 7, bitsCount: 3)
        // - 111: maximum repeat amount,
        writer.write(unsignedNumber: 7, bitsCount: 3)
        // - 00: encoded code length symbol 1 (which means code length is equal to 1),
        writer.write(unsignedNumber: 0, bitsCount: 2)
        // - 00: encoded code length symbol 1 (which means code length is equal to 1).
        writer.write(unsignedNumber: 0, bitsCount: 2)
        writer.align()

        // Assuming that there are no literals in the block we would still need both a length symbol from the
        // literals/length tree and a distance symbol from the distance tree to perform a match copy. Thus, this input
        // is corrupted.
        #expect(throws: (any Error).self) { try Deflate.decompress(data: writer.data) }

        // All code lengths are zero for the distance tree, and two non-zero code lengths for the literals/length tree.
        writer = LsbBitWriter()
        // - 101: last block bit and dynamical Huffman block type,
        writer.write(bits: [1, 0, 1])
        // - 11101: maximal amount of literal codes (29 + 257 = 286),
        writer.write(unsignedNumber: 29, bitsCount: 5)
        // - 00001: number of distances codes (1 + 1 = 2),
        writer.write(unsignedNumber: 1, bitsCount: 5)
        // - 1110: number of code length codes (4 + 14 = 18),
        writer.write(unsignedNumber: 14, bitsCount: 4)
        // - 000 011 010 000 000 000 000 000 000 000 000 000 000 010 000 011 000 010: code lenghts for code length codes,
        writer.write(unsignedNumber: 0, bitsCount: 3)
        writer.write(unsignedNumber: 3, bitsCount: 3)
        writer.write(unsignedNumber: 2, bitsCount: 3)
        for _ in 0..<10 {
            writer.write(unsignedNumber: 0, bitsCount: 3)
        }
        writer.write(unsignedNumber: 2, bitsCount: 3)
        writer.write(unsignedNumber: 0, bitsCount: 3)
        writer.write(unsignedNumber: 3, bitsCount: 3)
        writer.write(unsignedNumber: 0, bitsCount: 3)
        writer.write(unsignedNumber: 2, bitsCount: 3)
        // - 011: encoded code length symbol 2 (which means code length is equal to 2),
        writer.write(unsignedNumber: 3, bitsCount: 3)
        // - 01: encoded code length symbol 18 (which means to repeat 0 code length),
        writer.write(unsignedNumber: 1, bitsCount: 2)
        // - 1111111: maximum repeat amount,
        writer.write(unsignedNumber: 127, bitsCount: 7)
        // - 01: encoded code length symbol 18 (which means to repeat 0 code length),
        writer.write(unsignedNumber: 1, bitsCount: 2)
        // - 1101010: repeat amount,
        writer.write(unsignedNumber: 106, bitsCount: 7)
        // - 011: encoded code length symbol 3 (which means code length is equal to 3),
        writer.write(unsignedNumber: 2, bitsCount: 2)
        // - 01: encoded code length symbol 18 (which means to repeat 0 code length),
        writer.write(unsignedNumber: 1, bitsCount: 2)
        // - 0010100: repeat amount,
        writer.write(unsignedNumber: 20, bitsCount: 7)
        // - 00: literal symbol 0,
        writer.write(unsignedNumber: 0, bitsCount: 2)
        // - 010: EOS symbol 256,
        writer.write(unsignedNumber: 2, bitsCount: 3)
        writer.align()

        // The distance tree can be empty if a block contains only literals. In this case we encoded a single byte of 0.
        #expect(try Deflate.decompress(data: writer.data) == Data([0]))
    }

}