#include <unity.h>

#include <cstring>
#include <string>
#include <vector>

#include "png_reader.h"
#include "png_rgb332.h"
#include "pngs.h"

using namespace roro::files;

void setUp() {}
void tearDown() {}

namespace {
ImageRead from(const std::string& bytes) {
    return [&bytes](uint32_t at, uint8_t* into, size_t len) -> size_t {
        if (at >= bytes.size()) return 0;
        size_t n = std::min(len, bytes.size() - at);
        std::memcpy(into, bytes.data() + at, n);
        return n;
    };
}
template <size_t N>
std::string str(const uint8_t (&a)[N]) { return std::string(reinterpret_cast<const char*>(a), N); }

struct Picture {
    int w, h;
    std::vector<int> rgb;  // -1: never drawn
    Picture(int width, int height) : w(width), h(height), rgb(static_cast<size_t>(width * height), -1) {}
    ImagePixels sink() {
        return [this](int x, int y, int count, const uint8_t* p) {
            for (int i = 0; i < count; i++) {
                TEST_ASSERT_TRUE(x + i >= 0 && x + i < w && y >= 0 && y < h);
                rgb[static_cast<size_t>(y * w + x + i)] = (p[i * 3] << 16) | (p[i * 3 + 1] << 8) | p[i * 3 + 2];
            }
        };
    }
    int at(int x, int y) const { return rgb[static_cast<size_t>(y * w + x)]; }
};

Picture decode(const std::string& file, int w, int h) {
    Picture p(w, h);
    TEST_ASSERT_EQUAL_STRING("", readPng(from(file), static_cast<uint32_t>(file.size()), p.sink()).c_str());
    return p;
}
int colour(int x, int y) { return (((x * 9 + y) % 256) << 16) | (((y * 17) % 256) << 8) | ((x * x + y * y) % 256); }
int grey(int v) { return (v << 16) | (v << 8) | v; }
}  // namespace

void test_rgb() {
    Picture p = decode(str(kPngRgb), 96, 64);
    for (int y = 0; y < 64; y++)
        for (int x = 0; x < 96; x++) TEST_ASSERT_EQUAL_HEX32(colour(x, y), p.at(x, y));
}

void test_a_noisy_picture_with_every_kind_of_row() {
    Picture p = decode(str(kPngPhoto), 64, 48);
    for (int i = 0; i < 64 * 48; i++)
        TEST_ASSERT_EQUAL_HEX32((kPngPhotoPixels[i * 3] << 16) | (kPngPhotoPixels[i * 3 + 1] << 8) | kPngPhotoPixels[i * 3 + 2], p.rgb[static_cast<size_t>(i)]);
}

void test_stored_not_compressed() {
    Picture p = decode(str(kPngStored), 16, 12);
    for (int y = 0; y < 12; y++)
        for (int x = 0; x < 16; x++) TEST_ASSERT_EQUAL_HEX32(colour(x, y), p.at(x, y));
}

void test_references_far_back() {
    Picture p = decode(str(kPngFarBack), 120, 90);
    uint32_t total = 0;
    for (int v : p.rgb) {
        TEST_ASSERT_TRUE(v >= 0);
        total += static_cast<uint32_t>((v >> 16) + ((v >> 8) & 255) + (v & 255));
    }
    TEST_ASSERT_EQUAL_UINT32((kPngFarBackSum[0] << 24) | (kPngFarBackSum[1] << 16) | (kPngFarBackSum[2] << 8) | kPngFarBackSum[3], total);
    for (int x = 0; x < 120; x++) TEST_ASSERT_EQUAL_HEX32(p.at(x, 0) + 1, p.at(x, 60));  // the same row, its blue one more
}

void test_transparent_pixels_are_left_out() {
    Picture p = decode(str(kPngRgba), 16, 12);
    for (int y = 0; y < 12; y++)
        for (int x = 0; x < 16; x++) TEST_ASSERT_EQUAL_HEX32((x + y) % 5 == 0 ? -1 : colour(x, y), p.at(x, y));
    Picture g = decode(str(kPngGrayAlpha), 16, 12);
    for (int y = 0; y < 12; y++)
        for (int x = 0; x < 16; x++) TEST_ASSERT_EQUAL_HEX32(x % 2 ? grey((x * 16 + y * 3) % 256) : -1, g.at(x, y));
}

void test_palettes() {
    auto entry = [](int i) { return (((i * 7) % 256) << 16) | (((255 - i * 3) % 256) << 8) | ((i * 13 + 40) % 256); };
    Picture p = decode(str(kPngPalette), 16, 12);
    for (int y = 0; y < 12; y++)
        for (int x = 0; x < 16; x++) {
            int i = (x * 3 + y * 5) % 7;
            TEST_ASSERT_EQUAL_HEX32(i == 3 ? -1 : entry(i), p.at(x, y));  // index 3 is the transparent one
        }
    Picture q = decode(str(kPngPalette4), 17, 5);  // four bits a pixel, a row that ends in half a byte
    for (int y = 0; y < 5; y++)
        for (int x = 0; x < 17; x++) TEST_ASSERT_EQUAL_HEX32(entry((x + y * 3) % 13), q.at(x, y));
}

void test_greys() {
    Picture g = decode(str(kPngGray), 16, 12);
    Picture one = decode(str(kPngGray1), 19, 7);
    Picture deep = decode(str(kPngGray16), 16, 12);
    for (int y = 0; y < 12; y++)
        for (int x = 0; x < 16; x++) {
            TEST_ASSERT_EQUAL_HEX32(grey((x * 16 + y * 3) % 256), g.at(x, y));
            TEST_ASSERT_EQUAL_HEX32(grey(((x * 4000 + y * 300) % 65536) >> 8), deep.at(x, y));
        }
    for (int y = 0; y < 7; y++)
        for (int x = 0; x < 19; x++) TEST_ASSERT_EQUAL_HEX32((x + y) % 2 ? 0xFFFFFF : 0, one.at(x, y));
}

void test_one_of_the_firmwares_own_screenshots() {
    std::string file;
    roro::png::Rgb332Writer writer(40, 9, [&](const uint8_t* d, size_t n) {
        file.append(reinterpret_cast<const char*>(d), n);
        return true;
    });
    writer.begin();
    std::vector<uint8_t> row(40);
    for (int y = 0; y < 9; y++) {
        for (int x = 0; x < 40; x++) row[static_cast<size_t>(x)] = static_cast<uint8_t>(x * 6 + y);
        writer.row(row.data());
    }
    writer.end();
    Picture p = decode(file, 40, 9);
    for (int y = 0; y < 9; y++)
        for (int x = 0; x < 40; x++) {
            int v = (x * 6 + y) & 0xFF;
            TEST_ASSERT_EQUAL_HEX32((((v >> 5) * 255 / 7) << 16) | ((((v >> 2) & 7) * 255 / 7) << 8) | ((v & 3) * 85), p.at(x, y));
        }
}

void test_rows_can_be_left_out_and_the_end_cut_off() {
    std::string file = str(kPngRgb);
    Picture p(96, 64);
    int reads = 0;
    ImageRead counted = [&](uint32_t at, uint8_t* into, size_t len) {
        reads++;
        return from(file)(at, into, len);
    };
    TEST_ASSERT_EQUAL_STRING("", readPng(counted, static_cast<uint32_t>(file.size()), p.sink(), [](int y) { return y % 2 == 0; }, [](int y) { return y >= 10; }).c_str());
    TEST_ASSERT_EQUAL_HEX32(colour(5, 8), p.at(5, 8));
    TEST_ASSERT_EQUAL_INT(-1, p.at(5, 9));   // left out
    TEST_ASSERT_EQUAL_INT(-1, p.at(5, 10));  // and nothing from the tenth on
    Picture all(96, 64);
    int allReads = 0;
    ImageRead countAll = [&](uint32_t at, uint8_t* into, size_t len) {
        allReads++;
        return from(file)(at, into, len);
    };
    readPng(countAll, static_cast<uint32_t>(file.size()), all.sink());
    TEST_ASSERT_TRUE(reads < allReads / 2);  // it stopped reading the file too
}

void test_what_it_refuses() {
    Picture p(96, 64);
    std::string interlaced = str(kPngRgb);
    interlaced[28] = 1;
    TEST_ASSERT_EQUAL_STRING("An interlaced PNG can't be shown", readPng(from(interlaced), static_cast<uint32_t>(interlaced.size()), p.sink()).c_str());
    std::string cut = str(kPngRgb).substr(0, 2000);
    TEST_ASSERT_EQUAL_STRING("This PNG is cut short", readPng(from(cut), static_cast<uint32_t>(cut.size()), p.sink()).c_str());
    std::string notPng = "GIF89a and then some more bytes that are not a PNG";
    TEST_ASSERT_EQUAL_STRING("This PNG is damaged", readPng(from(notPng), static_cast<uint32_t>(notPng.size()), p.sink()).c_str());
    // Any byte changed: an answer, never a crash or a pixel outside the picture (the sink checks).
    std::string file = str(kPngPhoto);
    for (size_t i = 8; i < file.size(); i += 7) {
        std::string bad = file;
        bad[i] = static_cast<char>(bad[i] ^ 0xA5);
        Picture q(64, 48);
        readPng(from(bad), static_cast<uint32_t>(bad.size()), q.sink());
    }
}

int main() {
    UNITY_BEGIN();
    RUN_TEST(test_rgb);
    RUN_TEST(test_a_noisy_picture_with_every_kind_of_row);
    RUN_TEST(test_stored_not_compressed);
    RUN_TEST(test_references_far_back);
    RUN_TEST(test_transparent_pixels_are_left_out);
    RUN_TEST(test_palettes);
    RUN_TEST(test_greys);
    RUN_TEST(test_one_of_the_firmwares_own_screenshots);
    RUN_TEST(test_rows_can_be_left_out_and_the_end_cut_off);
    RUN_TEST(test_what_it_refuses);
    return UNITY_END();
}