#include <unity.h>

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

#include "image_file.h"
#include "images.h"
#include "png_rgb332.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); }

// The palette and the two patterns the test files were drawn with.
void colour(int i, uint8_t* rgb) {
    rgb[0] = static_cast<uint8_t>((i * 7) % 256);
    rgb[1] = static_cast<uint8_t>((255 - i * 3) % 256);
    rgb[2] = static_cast<uint8_t>((i * 13 + 40) % 256);
}
int small(int x, int y) { return (x * 3 + y * 5) % 7; }
int big(int x, int y) { return (x * x + y * y * 3 + x * y) & 0xFF; }

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];
            }
        };
    }
};

void checkGif(const std::string& file, int w, int h, int (*index)(int, int), int transparent = -1) {
    Picture p(w, h);
    TEST_ASSERT_EQUAL_STRING("", readGif(from(file), static_cast<uint32_t>(file.size()), p.sink()).c_str());
    for (int y = 0; y < h; y++)
        for (int x = 0; x < w; x++) {
            uint8_t c[3];
            colour(index(x, y), c);
            int want = index(x, y) == transparent ? -1 : (c[0] << 16) | (c[1] << 8) | c[2];
            char msg[40];
            std::snprintf(msg, sizeof msg, "at %d,%d", x, y);
            TEST_ASSERT_EQUAL_HEX32_MESSAGE(want, p.rgb[static_cast<size_t>(y * w + x)], msg);
        }
}
}  // namespace

void test_kinds() {
    TEST_ASSERT_TRUE(imageKindOfName("Photo.JPG") == ImageKind::Jpeg);
    TEST_ASSERT_TRUE(imageKindOfName("a.jpeg") == ImageKind::Jpeg);
    TEST_ASSERT_TRUE(imageKindOfName("20261007-104357.png") == ImageKind::Png);
    TEST_ASSERT_TRUE(imageKindOfName("x.bmp") == ImageKind::Bmp);
    TEST_ASSERT_TRUE(imageKindOfName("x.gif") == ImageKind::Gif);
    TEST_ASSERT_TRUE(imageKindOfName("notes.txt") == ImageKind::None);
    const uint8_t jpeg[] = {0xFF, 0xD8, 0xFF, 0xE0};
    TEST_ASSERT_TRUE(imageKindOfBytes(jpeg, 4) == ImageKind::Jpeg);
    TEST_ASSERT_TRUE(imageKindOfBytes(kGifSmall, 6) == ImageKind::Gif);
    TEST_ASSERT_TRUE(imageKindOfBytes(kBmp24, 8) == ImageKind::Bmp);
    TEST_ASSERT_TRUE(imageKindOfBytes(reinterpret_cast<const uint8_t*>("Hello"), 5) == ImageKind::None);
}

void test_sizes() {
    ImageInfo info;
    std::string gif = str(kGifBig), bmp = str(kBmp24);
    TEST_ASSERT_EQUAL_STRING("", imageInfo(from(gif), static_cast<uint32_t>(gif.size()), info).c_str());
    TEST_ASSERT_TRUE(info.kind == ImageKind::Gif && info.width == 96 && info.height == 64);
    TEST_ASSERT_EQUAL_STRING("", imageInfo(from(bmp), static_cast<uint32_t>(bmp.size()), info).c_str());
    TEST_ASSERT_TRUE(info.kind == ImageKind::Bmp && info.width == 13 && info.height == 7);

    // A JPEG: start, an application segment, a table, then the frame header with the size.
    std::string jpeg("\xFF\xD8", 2);
    jpeg += std::string("\xFF\xE1\x00\x08" "Exif\0\0", 10);
    jpeg += std::string("\xFF\xDB\x00\x04\x00\x00", 6);
    std::string frame("\xFF\xC0\x00\x11\x08\x0B\xB8\x0F\xA0\x03", 10);  // 3000 high, 4000 wide
    std::string baseline = jpeg + frame + std::string(20, '\0');
    TEST_ASSERT_EQUAL_STRING("", imageInfo(from(baseline), static_cast<uint32_t>(baseline.size()), info).c_str());
    TEST_ASSERT_TRUE(info.kind == ImageKind::Jpeg && info.width == 4000 && info.height == 3000);
    frame[1] = '\xC2';
    std::string progressive = jpeg + frame + std::string(20, '\0');
    TEST_ASSERT_EQUAL_STRING("A progressive JPEG can't be shown", imageInfo(from(progressive), static_cast<uint32_t>(progressive.size()), info).c_str());
    std::string cut = jpeg;
    TEST_ASSERT_EQUAL_STRING("This JPEG is damaged", imageInfo(from(cut), static_cast<uint32_t>(cut.size()), info).c_str());

    std::string text = "Just some words, not a picture at all.";
    TEST_ASSERT_EQUAL_STRING("Not a picture this can show", imageInfo(from(text), static_cast<uint32_t>(text.size()), info).c_str());
}

void test_a_screenshot_is_read_without_decoding() {
    const int w = 240, h = 135;
    std::string file;
    roro::png::Rgb332Writer writer(w, h, [&](const uint8_t* d, size_t n) {
        file.append(reinterpret_cast<const char*>(d), n);
        return true;
    });
    TEST_ASSERT_TRUE(writer.begin());
    std::vector<uint8_t> row(w);
    for (int y = 0; y < h; y++) {
        for (int x = 0; x < w; x++) row[static_cast<size_t>(x)] = static_cast<uint8_t>(x * 3 + y * 7);
        TEST_ASSERT_TRUE(writer.row(row.data()));
    }
    TEST_ASSERT_TRUE(writer.end());
    ImageInfo info;
    TEST_ASSERT_EQUAL_STRING("", imageInfo(from(file), static_cast<uint32_t>(file.size()), info).c_str());
    TEST_ASSERT_TRUE(info.kind == ImageKind::Png && info.width == w && info.height == h);
    uint32_t at = screenshotPixelsAt(from(file), static_cast<uint32_t>(file.size()), w, h);
    TEST_ASSERT_TRUE(at > 0);
    for (int y : {0, 1, 77, 134})
        for (int x : {0, 5, 239})
            TEST_ASSERT_EQUAL_UINT8(static_cast<uint8_t>(x * 3 + y * 7), static_cast<uint8_t>(file[at + static_cast<uint32_t>(y * (w + 1) + x)]));
    // Another PNG isn't taken for one: a byte more, a compressed block.
    std::string longer = file + "x";
    TEST_ASSERT_EQUAL_UINT32(0, screenshotPixelsAt(from(longer), static_cast<uint32_t>(longer.size()), w, h));
    std::string other = file;
    other[8 + 25 + 780 + 10] = 0x05;
    TEST_ASSERT_EQUAL_UINT32(0, screenshotPixelsAt(from(other), static_cast<uint32_t>(other.size()), w, h));
}

void test_dithering_leaves_the_screens_own_colours_alone() {
    for (int v = 0; v < 256; v++) {
        uint8_t r = static_cast<uint8_t>((v >> 5) * 255 / 7), g = static_cast<uint8_t>(((v >> 2) & 7) * 255 / 7), b = static_cast<uint8_t>((v & 3) * 85);
        for (int y = 0; y < 4; y++)
            for (int x = 0; x < 4; x++) TEST_ASSERT_EQUAL_UINT8(v, rgb332Dithered(r, g, b, x, y));
    }
    // A grey between two levels comes out as a mix of both, in proportion.
    int high = 0;
    for (int y = 0; y < 4; y++)
        for (int x = 0; x < 4; x++) {
            int red = rgb332Dithered(54, 0, 0, x, y) >> 5;  // half way from 36 to 72
            TEST_ASSERT_TRUE(red == 1 || red == 2);
            high += red == 2;
        }
    TEST_ASSERT_TRUE(high >= 6 && high <= 10);
    TEST_ASSERT_EQUAL_UINT8(0xFF, rgb332Dithered(255, 255, 255, 1, 2));
    TEST_ASSERT_EQUAL_UINT8(0x00, rgb332Dithered(0, 0, 0, 3, 3));
}

void test_a_small_picture_sits_in_the_middle() {
    ImageFrame f(100, 50, 0, 12, 240, 123);
    TEST_ASSERT_FALSE(f.bigger());
    TEST_ASSERT_EQUAL_INT(100, f.percent());
    ImageMap m = f.map();
    int tx, ty;
    TEST_ASSERT_TRUE(m.at(0, 0, tx, ty));
    TEST_ASSERT_EQUAL_INT(70, tx);
    TEST_ASSERT_EQUAL_INT(12 + 36, ty);
    TEST_ASSERT_TRUE(m.at(99, 49, tx, ty));
    TEST_ASSERT_EQUAL_INT(169, tx);
    TEST_ASSERT_FALSE(m.at(-1, 0, tx, ty));
    f.toggle();  // nothing to zoom
    TEST_ASSERT_FALSE(f.actual());
    TEST_ASSERT_EQUAL_INT(0, f.jpegShrink());
}

void test_a_big_picture_is_shrunk_to_fit_each_screen_pixel_once() {
    ImageFrame f(4000, 3000, 0, 12, 240, 123);
    TEST_ASSERT_TRUE(f.bigger());
    TEST_ASSERT_EQUAL_INT(4, f.percent());  // 123 / 3000
    ImageMap m = f.map();
    std::vector<int> hits(240 * 135, 0);
    int rows = 0;
    for (int y = 0; y < 3000; y++) {
        if (!m.rowUsed(y)) continue;
        rows++;
        for (int x = 0; x < 4000; x++) {
            int tx, ty;
            if (m.at(x, y, tx, ty)) hits[static_cast<size_t>(ty * 240 + tx)]++;
        }
    }
    TEST_ASSERT_EQUAL_INT(123, rows);
    int drawn = 0;
    for (int ty = 0; ty < 135; ty++)
        for (int tx = 0; tx < 240; tx++) {
            int n = hits[static_cast<size_t>(ty * 240 + tx)];
            TEST_ASSERT_TRUE(n <= 1);
            drawn += n;
            if (n) TEST_ASSERT_TRUE(ty >= 12 && tx >= 38 && tx < 202);  // 164 wide, centred
        }
    TEST_ASSERT_EQUAL_INT(164 * 123, drawn);

    // A JPEG decoder may halve it three times first; the map then takes those pixels.
    TEST_ASSERT_EQUAL_INT(3, f.jpegShrink());
    ImageMap j = f.map(3);
    int tx, ty, again = 0;
    for (int y = 0; y < 375; y++)
        for (int x = 0; x < 500; x++) again += j.at(x, y, tx, ty);
    TEST_ASSERT_EQUAL_INT(164 * 123, again);
    TEST_ASSERT_EQUAL_INT(0, ImageFrame(300, 200, 0, 12, 240, 123).jpegShrink());  // not by half: it would be too small
    TEST_ASSERT_EQUAL_INT(1, ImageFrame(480, 246, 0, 12, 240, 123).jpegShrink());
}

void test_at_its_own_size_it_moves_half_a_screen_at_a_time() {
    ImageFrame f(800, 600, 0, 12, 240, 123);
    f.toggle();
    TEST_ASSERT_TRUE(f.actual());
    TEST_ASSERT_EQUAL_INT(100, f.percent());
    ImageMap m = f.map();
    int tx, ty;
    TEST_ASSERT_TRUE(m.at(280, 238, tx, ty));  // the middle first
    TEST_ASSERT_EQUAL_INT(0, tx);
    TEST_ASSERT_EQUAL_INT(12, ty);
    TEST_ASSERT_FALSE(m.at(279, 238, tx, ty));
    TEST_ASSERT_FALSE(m.below(238 + 122));  // the last row on screen
    TEST_ASSERT_TRUE(m.below(238 + 123));   // from here on a decoder can stop
    TEST_ASSERT_TRUE(f.pan(1, 0));
    TEST_ASSERT_TRUE(f.map().at(400, 238, tx, ty));
    TEST_ASSERT_EQUAL_INT(0, tx);
    for (int i = 0; i < 10; i++) f.pan(1, 1);
    TEST_ASSERT_FALSE(f.pan(1, 0));  // the corner: no further
    TEST_ASSERT_TRUE(f.map().at(799, 599, tx, ty));
    TEST_ASSERT_EQUAL_INT(239, tx);
    TEST_ASSERT_EQUAL_INT(12 + 122, ty);
    f.toggle();
    TEST_ASSERT_FALSE(f.actual());
    TEST_ASSERT_FALSE(f.pan(1, 0));  // fitted: nothing to move
}

void test_gif() {
    checkGif(str(kGifSmall), 16, 12, small);
    checkGif(str(kGifSmallInterlaced), 16, 12, small);
    checkGif(str(kGifBig), 96, 64, big);  // long enough for the codes to grow to 12 bits and start over
    checkGif(str(kGifBigInterlaced), 96, 64, big);
    checkGif(str(kGifTransparent), 16, 12, small, 3);
}

void test_gif_damaged() {
    Picture p(96, 64);
    std::string cut = str(kGifBig).substr(0, 3000);
    TEST_ASSERT_EQUAL_STRING("This GIF is cut short", readGif(from(cut), static_cast<uint32_t>(cut.size()), p.sink()).c_str());
    std::string head = str(kGifBig).substr(0, 10);
    TEST_ASSERT_EQUAL_STRING("This GIF is damaged", readGif(from(head), static_cast<uint32_t>(head.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(kGifBig);
    for (size_t i = 0; i < file.size(); i += 37) {
        std::string bad = file;
        bad[i] = static_cast<char>(bad[i] ^ 0x5A);
        Picture q(96, 64);
        readGif(from(bad), static_cast<uint32_t>(bad.size()), q.sink());
    }
}

void test_bmp() {
    std::string file = str(kBmp24);
    Picture p(13, 7);
    TEST_ASSERT_EQUAL_STRING("", readBmp(from(file), static_cast<uint32_t>(file.size()), p.sink()).c_str());
    for (int y = 0; y < 7; y++)
        for (int x = 0; x < 13; x++)
            TEST_ASSERT_EQUAL_HEX32((((x * 9) % 256) << 16) | (((y * 17) % 256) << 8) | ((x + y) % 256), p.rgb[static_cast<size_t>(y * 13 + x)]);
    std::string file8 = str(kBmp8);
    Picture q(13, 7);
    TEST_ASSERT_EQUAL_STRING("", readBmp(from(file8), static_cast<uint32_t>(file8.size()), q.sink()).c_str());
    for (int i = 0; i < 13 * 7; i++)
        TEST_ASSERT_EQUAL_HEX32((kBmp8Pixels[i * 3] << 16) | (kBmp8Pixels[i * 3 + 1] << 8) | kBmp8Pixels[i * 3 + 2], q.rgb[static_cast<size_t>(i)]);
    // Only the rows asked for are read.
    Picture r(13, 7);
    readBmp(from(file), static_cast<uint32_t>(file.size()), r.sink(), [](int y) { return y == 2; });
    TEST_ASSERT_EQUAL_INT(-1, r.rgb[0]);
    TEST_ASSERT_TRUE(r.rgb[2 * 13] >= 0);
    std::string cut = file.substr(0, 200);
    TEST_ASSERT_EQUAL_STRING("This BMP is cut short", readBmp(from(cut), static_cast<uint32_t>(cut.size()), p.sink()).c_str());
}

int main() {
    UNITY_BEGIN();
    RUN_TEST(test_kinds);
    RUN_TEST(test_sizes);
    RUN_TEST(test_a_screenshot_is_read_without_decoding);
    RUN_TEST(test_dithering_leaves_the_screens_own_colours_alone);
    RUN_TEST(test_a_small_picture_sits_in_the_middle);
    RUN_TEST(test_a_big_picture_is_shrunk_to_fit_each_screen_pixel_once);
    RUN_TEST(test_at_its_own_size_it_moves_half_a_screen_at_a_time);
    RUN_TEST(test_gif);
    RUN_TEST(test_gif_damaged);
    RUN_TEST(test_bmp);
    return UNITY_END();
}