#include "image_file.h"

#include <algorithm>
#include <cstring>
#include <memory>
#include <new>

#include "file_names.h"
#include "png_rgb332.h"

namespace roro::files {

namespace {
uint32_t le16(const uint8_t* p) { return p[0] | (p[1] << 8); }
uint32_t le32(const uint8_t* p) { return p[0] | (p[1] << 8) | (p[2] << 16) | (static_cast<uint32_t>(p[3]) << 24); }
uint32_t be16(const uint8_t* p) { return (p[0] << 8) | p[1]; }
uint32_t be32(const uint8_t* p) { return (static_cast<uint32_t>(p[0]) << 24) | (p[1] << 16) | (p[2] << 8) | p[3]; }

constexpr int kMaxSide = 16384;  // more than that isn't a picture for this screen
const uint8_t kPngSignature[] = {0x89, 'P', 'N', 'G', '\r', '\n', 0x1A, '\n'};

// A file read from its start on, a small block at a time.
class Stream {
   public:
    Stream(const ImageRead& read, uint32_t size) : read_(read), size_(size) {}
    int get() {
        if (at_ >= have_) {
            if (next_ >= size_) return -1;
            have_ = read_(next_, buf_, std::min<size_t>(sizeof buf_, size_ - next_));
            at_ = 0;
            next_ += static_cast<uint32_t>(have_);
            if (!have_) return -1;
        }
        return buf_[at_++];
    }
    bool take(uint8_t* into, size_t len) {
        for (size_t i = 0; i < len; i++) {
            int c = get();
            if (c < 0) return false;
            into[i] = static_cast<uint8_t>(c);
        }
        return true;
    }
    bool skip(size_t len) {
        size_t buffered = std::min(len, have_ - at_);
        at_ += buffered;
        len -= buffered;
        if (len > size_ - next_) return false;
        next_ += static_cast<uint32_t>(len);
        return true;
    }

   private:
    const ImageRead& read_;
    uint32_t size_, next_ = 0;
    uint8_t buf_[256];
    size_t have_ = 0, at_ = 0;
};
}  // namespace

ImageKind imageKindOfName(const std::string& name) {
    std::string ext = extensionOf(name);
    if (ext == "png") return ImageKind::Png;
    if (ext == "jpg" || ext == "jpeg") return ImageKind::Jpeg;
    if (ext == "bmp") return ImageKind::Bmp;
    if (ext == "gif") return ImageKind::Gif;
    return ImageKind::None;
}

ImageKind imageKindOfBytes(const uint8_t* head, size_t len) {
    if (len >= 8 && std::memcmp(head, kPngSignature, 8) == 0) return ImageKind::Png;
    if (len >= 3 && head[0] == 0xFF && head[1] == 0xD8 && head[2] == 0xFF) return ImageKind::Jpeg;
    if (len >= 6 && (std::memcmp(head, "GIF87a", 6) == 0 || std::memcmp(head, "GIF89a", 6) == 0)) return ImageKind::Gif;
    if (len >= 2 && head[0] == 'B' && head[1] == 'M') return ImageKind::Bmp;
    return ImageKind::None;
}

const char* imageKindName(ImageKind kind) {
    switch (kind) {
        case ImageKind::Png: return "PNG";
        case ImageKind::Jpeg: return "JPEG";
        case ImageKind::Bmp: return "BMP";
        case ImageKind::Gif: return "GIF";
        default: return "";
    }
}

std::string imageInfo(const ImageRead& read, uint32_t size, ImageInfo& out) {
    uint8_t head[32];
    size_t n = read(0, head, std::min<size_t>(sizeof head, size));
    out.kind = imageKindOfBytes(head, n);
    long w = 0, h = 0;
    switch (out.kind) {
        case ImageKind::Png:
            if (n < 24 || std::memcmp(head + 12, "IHDR", 4) != 0) return "This PNG is damaged";
            w = static_cast<long>(be32(head + 16));
            h = static_cast<long>(be32(head + 20));
            break;
        case ImageKind::Gif:
            if (n < 10) return "This GIF is damaged";
            w = static_cast<long>(le16(head + 6));
            h = static_cast<long>(le16(head + 8));
            break;
        case ImageKind::Bmp: {
            if (n < 26) return "This BMP is damaged";
            uint32_t dib = le32(head + 14);
            if (dib < 40) return "This kind of BMP can't be shown";
            w = static_cast<int32_t>(le32(head + 18));
            h = static_cast<int32_t>(le32(head + 22));
            if (h < 0) h = -h;  // top row first
            break;
        }
        case ImageKind::Jpeg: {
            // Marker after marker until the one that carries the size.
            uint32_t at = 2;
            for (int guard = 0; guard < 4000; guard++) {
                uint8_t m[9];
                if (read(at, m, 4) != 4 || m[0] != 0xFF) return "This JPEG is damaged";
                uint8_t marker = m[1];
                if (marker == 0xFF) {  // padding
                    at++;
                    continue;
                }
                if (marker == 0x01 || (marker >= 0xD0 && marker <= 0xD8)) {  // no length
                    at += 2;
                    continue;
                }
                if (marker == 0xD9 || marker == 0xDA) return "This JPEG is damaged";  // the picture, and no size yet
                bool frame = marker >= 0xC0 && marker <= 0xCF && marker != 0xC4 && marker != 0xC8 && marker != 0xCC;
                if (frame) {
                    if (read(at, m, 9) != 9) return "This JPEG is damaged";
                    h = static_cast<long>(be16(m + 5));
                    w = static_cast<long>(be16(m + 7));
                    if (marker == 0xC2) return "A progressive JPEG can't be shown";
                    if (marker != 0xC0 && marker != 0xC1) return "This kind of JPEG can't be shown";
                    break;
                }
                at += 2 + be16(m + 2);
            }
            break;
        }
        default: return "Not a picture this can show";
    }
    if (w <= 0 || h <= 0) return "This picture is damaged";
    if (w > kMaxSide || h > kMaxSide) return "Too big: 16,384 pixels a side at most";
    out.width = static_cast<int>(w);
    out.height = static_cast<int>(h);
    return "";
}

uint32_t screenshotPixelsAt(const ImageRead& read, uint32_t size, int width, int height) {
    if (width <= 0 || height <= 0 || size != png::Rgb332Writer::fileSize(width, height)) return 0;
    // Signature, IHDR, then a palette of 256 colours, then the one IDAT: a zlib header and a
    // single stored block.
    uint8_t ihdr[5], plte[8], idat[15];
    constexpr uint32_t kPlteAt = 8 + 25, kIdatAt = kPlteAt + 12 + 768;
    if (read(24, ihdr, 5) != 5 || ihdr[0] != 8 || ihdr[1] != 3 || ihdr[4] != 0) return 0;
    if (read(kPlteAt, plte, 8) != 8 || be32(plte) != 768 || std::memcmp(plte + 4, "PLTE", 4) != 0) return 0;
    if (read(kIdatAt, idat, 15) != 15 || std::memcmp(idat + 4, "IDAT", 4) != 0) return 0;
    uint32_t raw = static_cast<uint32_t>(width + 1) * static_cast<uint32_t>(height);
    if (idat[8] != 0x78 || idat[10] != 0x01 || le16(idat + 11) != raw || le16(idat + 13) != (raw ^ 0xFFFF)) return 0;
    return kIdatAt + 15 + 1;  // past the first row's filter byte
}

uint8_t rgb332Dithered(uint8_t r, uint8_t g, uint8_t b, int x, int y) {
    static const uint8_t kBayer[16] = {0, 8, 2, 10, 12, 4, 14, 6, 3, 11, 1, 9, 15, 7, 13, 5};
    int threshold = kBayer[((y & 3) << 2) | (x & 3)] * 16 + 8;  // 8 to 248
    auto level = [threshold](int v, int top) {
        int nearest = (v * top + 127) / 255;
        if (nearest * 255 / top == v) return nearest;  // a colour the screen has
        int low = v * top / 255, rest = v * top - low * 255;
        return rest > threshold ? low + 1 : low;
    };
    return static_cast<uint8_t>((level(r, 7) << 5) | (level(g, 7) << 2) | level(b, 3));
}

bool ImageMap::at(int sx, int sy, int& tx, int& ty) const {
    if (sx < 0 || sy < 0) return false;
    if (scale >= 65536) {
        tx = sx + offX;
        ty = sy + offY;
    } else {
        uint64_t fx = static_cast<uint64_t>(sx) * scale, fy = static_cast<uint64_t>(sy) * scale;
        if ((fx & 0xFFFF) >= scale || (fy & 0xFFFF) >= scale) return false;
        tx = static_cast<int>(fx >> 16) + offX;
        ty = static_cast<int>(fy >> 16) + offY;
    }
    if (tx < 0 || ty < 0 || tx >= viewW || ty >= viewH) return false;
    tx += viewX;
    ty += viewY;
    return true;
}

bool ImageMap::rowUsed(int sy) const {
    if (sy < 0) return false;
    int ty;
    if (scale >= 65536) {
        ty = sy + offY;
    } else {
        uint64_t fy = static_cast<uint64_t>(sy) * scale;
        if ((fy & 0xFFFF) >= scale) return false;
        ty = static_cast<int>(fy >> 16) + offY;
    }
    return ty >= 0 && ty < viewH;
}

bool ImageMap::below(int sy) const {
    if (sy < 0) return false;
    int ty = scale >= 65536 ? sy + offY : static_cast<int>((static_cast<uint64_t>(sy) * scale) >> 16) + offY;
    return ty >= viewH;
}

ImageFrame::ImageFrame(int width, int height, int viewX, int viewY, int viewW, int viewH)
    : w_(std::max(1, width)), h_(std::max(1, height)), vx_(viewX), vy_(viewY), vw_(std::max(1, viewW)), vh_(std::max(1, viewH)) {}

uint32_t ImageFrame::fitScale() const {
    uint64_t sx = (static_cast<uint64_t>(vw_) << 16) / static_cast<uint64_t>(w_), sy = (static_cast<uint64_t>(vh_) << 16) / static_cast<uint64_t>(h_);
    return static_cast<uint32_t>(std::min<uint64_t>(65536, std::max<uint64_t>(1, std::min(sx, sy))));
}

void ImageFrame::toggle() {
    if (!bigger()) return;
    actual_ = !actual_;
    if (actual_) {  // the middle of it first
        panX_ = std::max(0, (w_ - vw_) / 2);
        panY_ = std::max(0, (h_ - vh_) / 2);
    }
}

bool ImageFrame::pan(int dx, int dy) {
    if (!actual_) return false;
    int x = std::clamp(panX_ + dx * (vw_ / 2), 0, std::max(0, w_ - vw_));
    int y = std::clamp(panY_ + dy * (vh_ / 2), 0, std::max(0, h_ - vh_));
    bool moved = x != panX_ || y != panY_;
    panX_ = x;
    panY_ = y;
    return moved;
}

int ImageFrame::percent() const { return actual_ ? 100 : static_cast<int>((static_cast<uint64_t>(fitScale()) * 100 + 32768) >> 16); }

int ImageFrame::jpegShrink() const {
    if (actual_) return 0;
    uint32_t scale = fitScale();
    int shrink = 0;
    while (shrink < 3 && (static_cast<uint64_t>(scale) << (shrink + 1)) <= 65536) shrink++;
    return shrink;
}

ImageMap ImageFrame::map(int shrink) const {
    ImageMap m;
    m.viewX = vx_;
    m.viewY = vy_;
    m.viewW = vw_;
    m.viewH = vh_;
    if (actual_) {
        m.scale = 65536;
        m.offX = w_ <= vw_ ? (vw_ - w_) / 2 : -panX_;
        m.offY = h_ <= vh_ ? (vh_ - h_) / 2 : -panY_;
        return m;
    }
    uint32_t scale = fitScale();
    int tw = std::max<int>(1, static_cast<int>((static_cast<uint64_t>(w_) * scale) >> 16));
    int th = std::max<int>(1, static_cast<int>((static_cast<uint64_t>(h_) * scale) >> 16));
    m.offX = (vw_ - tw) / 2;
    m.offY = (vh_ - th) / 2;
    m.scale = static_cast<uint32_t>(std::min<uint64_t>(65536, static_cast<uint64_t>(scale) << shrink));
    return m;
}

std::string readBmp(const ImageRead& read, uint32_t size, const ImagePixels& pixels, const std::function<bool(int y)>& rowNeeded) {
    uint8_t head[54];
    if (read(0, head, sizeof head) != sizeof head || head[0] != 'B' || head[1] != 'M') return "This BMP is damaged";
    uint32_t dataAt = le32(head + 10), dib = le32(head + 14), compression = le32(head + 30), colours = le32(head + 46);
    int32_t w = static_cast<int32_t>(le32(head + 18)), h = static_cast<int32_t>(le32(head + 22));
    uint32_t bits = le16(head + 28);
    bool topDown = h < 0;
    if (topDown) h = -h;
    if (dib < 40 || w <= 0 || h <= 0 || w > kMaxSide || h > kMaxSide) return "This kind of BMP can't be shown";
    if ((bits != 8 && bits != 24 && bits != 32) || (compression != 0 && !(compression == 3 && bits == 32))) return "This kind of BMP can't be shown";
    std::unique_ptr<uint8_t[]> palette;
    if (bits == 8) {
        if (!colours || colours > 256) colours = 256;
        palette.reset(new (std::nothrow) uint8_t[1024]());
        if (!palette) return "Not enough memory";
        if (read(14 + dib, palette.get(), colours * 4) != colours * 4) return "This BMP is damaged";
    }
    uint32_t bytes = bits / 8, rowSize = (static_cast<uint32_t>(w) * bytes + 3) & ~3u;
    if (static_cast<uint64_t>(dataAt) + static_cast<uint64_t>(rowSize) * static_cast<uint32_t>(h) > size) return "This BMP is cut short";
    // A row in one read when it fits, a piece of it at a time otherwise.
    constexpr int kOut = 64;                  // pixels handed on at once
    constexpr uint32_t kRowBuffer = 4096;     // bytes
    uint32_t rowBytes = static_cast<uint32_t>(w) * bytes, bufSize = std::min(rowBytes, kRowBuffer);
    bufSize -= bufSize % bytes;
    std::unique_ptr<uint8_t[]> in(new (std::nothrow) uint8_t[bufSize]);
    if (!in) return "Not enough memory";
    uint8_t out[kOut * 3];
    // In the order the file has them, which is usually the last row first: going back through a
    // file on the card costs far more than going on (measured: a second for 135 rows).
    for (int stored = 0; stored < h; stored++) {
        int y = topDown ? stored : h - 1 - stored;
        if (rowNeeded && !rowNeeded(y)) continue;
        uint32_t rowAt = dataAt + rowSize * static_cast<uint32_t>(stored);
        for (uint32_t done = 0; done < rowBytes; done += bufSize) {
            size_t want = std::min(bufSize, rowBytes - done);
            if (read(rowAt + done, in.get(), want) != want) return "The card refused to read it";
            int first = static_cast<int>(done / bytes), count = static_cast<int>(want / bytes);
            for (int at = 0; at < count; at += kOut) {
                int n = std::min(kOut, count - at);
                for (int i = 0; i < n; i++) {
                    const uint8_t* p = bits == 8 ? palette.get() + in[at + i] * 4 : in.get() + static_cast<size_t>(at + i) * bytes;
                    out[i * 3] = p[2];  // stored blue, green, red
                    out[i * 3 + 1] = p[1];
                    out[i * 3 + 2] = p[0];
                }
                pixels(first + at, y, n, out);
            }
        }
    }
    return "";
}

namespace {
struct GifWork {
    uint8_t palette[768];
    uint16_t prefix[4096];
    uint8_t suffix[4096], stack[4096];
};
}  // namespace

std::string readGif(const ImageRead& read, uint32_t size, const ImagePixels& pixels) {
    Stream in(read, size);
    uint8_t head[13];
    if (!in.take(head, 13) || imageKindOfBytes(head, 6) != ImageKind::Gif) return "This GIF is damaged";
    int screenW = static_cast<int>(le16(head + 6)), screenH = static_cast<int>(le16(head + 8));
    std::unique_ptr<GifWork> work(new (std::nothrow) GifWork);
    if (!work) return "Not enough memory to show a GIF";
    std::memset(work->palette, 0, sizeof work->palette);
    if (head[10] & 0x80 && !in.take(work->palette, 3u << ((head[10] & 7) + 1))) return "This GIF is damaged";
    int transparent = -1;
    for (int guard = 0; guard < 100000; guard++) {
        int kind = in.get();
        if (kind == 0x21) {  // an extension: only the one before a picture matters, for its transparent colour
            int label = in.get();
            for (bool first = true;; first = false) {
                int len = in.get();
                if (len < 0) return "This GIF is damaged";
                if (len == 0) break;
                uint8_t block[255];
                if (!in.take(block, static_cast<size_t>(len))) return "This GIF is damaged";
                if (label == 0xF9 && first && len >= 4) transparent = (block[0] & 1) ? block[3] : -1;
            }
            continue;
        }
        if (kind != 0x2C) return kind == 0x3B ? "This GIF has no picture" : "This GIF is damaged";
        break;
    }
    uint8_t desc[9];
    if (!in.take(desc, 9)) return "This GIF is damaged";
    int left = static_cast<int>(le16(desc)), top = static_cast<int>(le16(desc + 2));
    int fw = static_cast<int>(le16(desc + 4)), fh = static_cast<int>(le16(desc + 6));
    bool interlaced = desc[8] & 0x40;
    if (desc[8] & 0x80 && !in.take(work->palette, 3u << ((desc[8] & 7) + 1))) return "This GIF is damaged";
    int minBits = in.get();
    if (fw <= 0 || fh <= 0 || minBits < 2 || minBits > 8) return "This GIF is damaged";

    // The pixels come out in the order they are stored; an interlaced picture stores every eighth
    // row first, then the rows between, in four passes.
    static const int kStart[4] = {0, 4, 2, 1}, kStep[4] = {8, 8, 4, 2};
    int px = 0, row = 0, pass = 0, rowsDone = 0;
    uint8_t run[64 * 3];
    int runLen = 0, runX = 0;
    auto flush = [&]() {
        int y = top + row;
        if (runLen && y >= 0 && y < screenH) pixels(left + runX, y, runLen, run);
        runLen = 0;
    };
    auto put = [&](uint8_t index) {
        if (rowsDone >= fh) return;
        int x = left + px;
        if (index == transparent || x < 0 || x >= screenW) {
            flush();
        } else {
            if (!runLen) runX = px;
            std::memcpy(run + runLen * 3, work->palette + index * 3, 3);
            if (++runLen == 64) flush();
        }
        if (++px < fw) return;
        flush();
        px = 0;
        rowsDone++;
        if (!interlaced) {
            row++;
            return;
        }
        row += kStep[pass];
        while (row >= fh && pass < 3) row = kStart[++pass];
    };

    const int clear = 1 << minBits, stop = clear + 1;
    int bits = minBits + 1, next = clear + 2, prev = -1, first = 0;
    uint32_t hold = 0;
    int held = 0, blockLeft = 0;
    bool ended = false;
    for (int i = 0; i < clear; i++) work->suffix[i] = static_cast<uint8_t>(i);
    while (rowsDone < fh && !ended) {
        while (held < bits) {
            if (!blockLeft) {
                blockLeft = in.get();
                if (blockLeft <= 0) {
                    ended = true;
                    break;
                }
            }
            int c = in.get();
            if (c < 0) return "This GIF is cut short";
            blockLeft--;
            hold |= static_cast<uint32_t>(c) << held;
            held += 8;
        }
        if (ended) break;
        int code = static_cast<int>(hold & ((1u << bits) - 1));
        hold >>= bits;
        held -= bits;
        if (code == clear) {
            bits = minBits + 1;
            next = clear + 2;
            prev = -1;
            continue;
        }
        if (code == stop) break;
        if (prev < 0) {
            if (code >= clear) return "This GIF is damaged";
            put(static_cast<uint8_t>(code));
            first = prev = code;
            continue;
        }
        if (code > next) return "This GIF is damaged";
        int sp = 0, walk = code;
        if (code == next) {  // the string being defined: the one before, and its own first pixel again
            work->stack[sp++] = static_cast<uint8_t>(first);
            walk = prev;
        }
        while (walk >= clear && sp < 4095) {
            work->stack[sp++] = work->suffix[walk];
            walk = work->prefix[walk];
        }
        if (walk >= clear) return "This GIF is damaged";
        first = walk;
        work->stack[sp++] = static_cast<uint8_t>(walk);
        if (next < 4096) {
            work->prefix[next] = static_cast<uint16_t>(prev);
            work->suffix[next] = static_cast<uint8_t>(first);
            next++;
            if (next == (1 << bits) && bits < 12) bits++;
        }
        prev = code;
        while (sp) put(work->stack[--sp]);
    }
    flush();
    return rowsDone ? "" : "This GIF is damaged";
}

}  // namespace roro::files