#include "image_pane.h"

#include <Arduino.h>
#include <SD.h>

#include <atomic>
#include <cstring>
#include <memory>
#include <new>

#include <lgfx/utility/lgfx_tjpgd.h>

#include "app_keys.h"
#include "file_names.h"
#include "platform/console.h"
#include "png_reader.h"
#include "ui/fonts.h"
#include "ui/theme.h"

namespace roro {

namespace {
constexpr uint32_t kNoteMs = 3000;
constexpr uint32_t kPushMs = 250;  // how often the screen shows how far the decoding is
constexpr int kNoteHeight = 11;
constexpr size_t kJpegPool = 3900;  // what the library gives its own JPEG decoder
}  // namespace

struct ImagePane::Job {
    std::string path, why;
    files::ImageInfo info;
    files::ImageFrame frame;
    files::ImageMap map;
    uint32_t size = 0, shotAt = 0, tookMs = 0;
    uint8_t* screen = nullptr;  // the screen's buffer: one byte a pixel, RRRGGGBB
    int stride = 0;
    std::atomic<bool> stop{false}, done{false};
    bool enough = false;  // the rest of the file is under the view: the decoder is told to stop
    File* file = nullptr;
    uint32_t sinceRest = 0;

    void put(int sx, int sy, uint8_t r, uint8_t g, uint8_t b) {
        int tx, ty;
        if (map.at(sx, sy, tx, ty)) screen[ty * stride + tx] = files::rgb332Dithered(r, g, b, tx, ty);
    }
    // A long decoding must leave the processor to others now and then (the idle task is watched).
    void rest(uint32_t bytes) {
        sinceRest += bytes;
        if (sinceRest < 16 * 1024) return;
        sinceRest = 0;
        vTaskDelay(1);
    }
    size_t readAt(uint32_t at, uint8_t* into, size_t len) {
        if (stop || !file->seek(at)) return 0;
        int n = file->read(into, len);
        rest(static_cast<uint32_t>(len));
        return n > 0 ? static_cast<size_t>(n) : 0;
    }
    void run();
};

namespace {
// The library's JPEG decoder reads the file from its start on; a null buffer means "skip".
// Nothing more to read is how a decoding is told to stop.
uint32_t readNext(void* job, uint8_t* into, uint32_t len) {
    auto& j = *static_cast<ImagePane::Job*>(job);
    if (j.stop || j.enough) return 0;
    File& f = *j.file;
    j.rest(len);
    if (!into) return f.seek(f.position() + len) ? len : 0;
    int n = f.read(into, len);
    return n > 0 ? static_cast<uint32_t>(n) : 0;
}

// A block of a JPEG: its pixels row by row, three bytes each.
uint32_t jpegBlock(void* job, void* bitmap, JRECT* rect) {
    auto& j = *static_cast<ImagePane::Job*>(job);
    const uint8_t* p = static_cast<const uint8_t*>(bitmap);
    if (j.map.below(static_cast<int>(rect->top))) return j.enough = true, 0;
    for (uint32_t y = rect->top; y <= rect->bottom; y++)
        for (uint32_t x = rect->left; x <= rect->right; x++, p += 3) j.put(static_cast<int>(x), static_cast<int>(y), p[0], p[1], p[2]);
    return j.stop ? 0 : 1;
}
}  // namespace

// On the storage task.
void ImagePane::Job::run() {
    uint32_t started = millis();
    File f = SD.open(path.c_str(), FILE_READ);
    if (!f) {
        why = "The card refused to open it";
        done = true;
        return;
    }
    file = &f;
    files::ImageRead read = [this](uint32_t at, uint8_t* into, size_t len) { return readAt(at, into, len); };
    files::ImagePixels pixels = [this](int x, int y, int count, const uint8_t* rgb) {
        for (int i = 0; i < count; i++, rgb += 3) put(x + i, y, rgb[0], rgb[1], rgb[2]);
    };
    switch (info.kind) {
        case files::ImageKind::Png:
            if (shotAt) {  // one of ours: each byte is already a colour of the screen
                std::unique_ptr<uint8_t[]> row(new (std::nothrow) uint8_t[info.width]);
                if (!row) {
                    why = "Not enough memory";
                    break;
                }
                for (int y = 0; y < info.height && why.empty() && !stop; y++) {
                    if (!map.rowUsed(y)) continue;
                    size_t w = static_cast<size_t>(info.width);
                    if (readAt(shotAt + static_cast<uint32_t>(y) * (info.width + 1), row.get(), w) != w) why = "The card refused to read it";
                    int tx, ty;
                    for (int x = 0; x < info.width; x++)
                        if (map.at(x, y, tx, ty)) screen[ty * stride + tx] = row[x];
                }
                break;
            }
            why = files::readPng(read, size, pixels, [this](int y) { return map.rowUsed(y); }, [this](int y) { return map.below(y); });
            break;
        case files::ImageKind::Jpeg: {
            std::unique_ptr<lgfxJdec> jpeg(new (std::nothrow) lgfxJdec);
            std::unique_ptr<uint8_t[]> pool(new (std::nothrow) uint8_t[kJpegPool]);
            if (!jpeg || !pool) {
                why = "Not enough memory";
                break;
            }
            int shrink = frame.jpegShrink();
            map = frame.map(shrink);
            JRESULT r = lgfx_jd_prepare(jpeg.get(), readNext, pool.get(), kJpegPool, this);
            if (r == JDR_OK) r = lgfx_jd_decomp(jpeg.get(), jpegBlock, static_cast<uint_fast8_t>(shrink));
            if (r == JDR_FMT3) why = "This kind of JPEG can't be shown";
            else if (r == JDR_MEM1 || r == JDR_MEM2) why = "This JPEG is too complex to show";
            else if (r != JDR_OK && !enough) why = "This JPEG is damaged";
            break;
        }
        case files::ImageKind::Bmp:
            why = files::readBmp(read, size, pixels, [this](int y) { return map.rowUsed(y); });
            break;
        case files::ImageKind::Gif: why = files::readGif(read, size, pixels); break;
        default: why = "Not a picture this can show"; break;
    }
    f.close();
    file = nullptr;
    tookMs = millis() - started;
    if (stop) why.clear();
    else
        console.printf("image: %s, %d x %d %s, %s in %u ms\n", path.c_str(), info.width, info.height, files::imageKindName(info.kind),
                       why.empty() ? (frame.actual() ? "its own size" : "fitted") : why.c_str(), (unsigned)tookMs);
    done = true;
}

std::string ImagePane::open(const std::string& path, uint32_t size) {
    close();
    std::string why;
    files::ImageInfo info;
    uint32_t shotAt = 0;
    bool ran = storage_.runAndWait([&]() {
        File f = SD.open(path.c_str(), FILE_READ);
        if (!f) {
            why = "The card refused to open it";
            return;
        }
        files::ImageRead read = [&f](uint32_t at, uint8_t* into, size_t len) -> size_t {
            if (!f.seek(at)) return 0;
            int n = f.read(into, len);
            return n > 0 ? static_cast<size_t>(n) : 0;
        };
        why = files::imageInfo(read, size, info);
        if (why.empty() && info.kind == files::ImageKind::Png) shotAt = files::screenshotPixelsAt(read, size, info.width, info.height);
        f.close();
    });
    if (!ran) why = "No SD card";
    if (!why.empty()) return why;
    path_ = path;
    size_ = size;
    info_ = info;
    shotAt_ = shotAt;
    const auto& area = theme::kContent;
    frame_ = files::ImageFrame(info.width, info.height, area.x, area.y, area.w, area.h);
    phase_ = Phase::Wanted;
    told_ = noteDrawn_ = false;
    problem_.clear();
    note_.clear();
    return "";
}

void ImagePane::cancel() {
    if (job_ && !job_->done) {
        job_->stop = true;
        storage_.runAndWait([]() {});  // behind the decoding in the queue: back when it has ended
    }
    job_.reset();
}

void ImagePane::close() {
    cancel();
    path_.clear();
    problem_.clear();
    note_.clear();
    info_ = files::ImageInfo();
    phase_ = Phase::Wanted;
    std::vector<uint8_t>().swap(under_);
}

void ImagePane::lost() {
    cancel();
    phase_ = Phase::Wanted;
    noteDrawn_ = false;
}

std::string ImagePane::describe() const {
    std::string s = std::to_string(info_.width) + " x " + std::to_string(info_.height) + " " + files::imageKindName(info_.kind);
    if (frame_.bigger()) s += frame_.actual() ? ", its own size" : ", at " + std::to_string(frame_.percent()) + " %";
    return s;
}

void ImagePane::say(const std::string& text) {
    hideNote(canvas_);
    note_ = text;
    noteMs_ = millis();
}

void ImagePane::help(std::vector<KeyHelp>& out) const { keys::add(out, keys::kViewerImage); }

bool ImagePane::onKey(const KeyEvent& e) {
    if (path_.empty()) return false;
    bool changed = false;
    switch (e.key) {
        case Key::Select:
            if (!frame_.bigger()) return true;
            cancel();  // before the frame it is drawing into changes
            frame_.toggle();
            told_ = false;
            changed = true;
            break;
        case Key::Up:
        case Key::Down:
        case Key::Left:
        case Key::Right: {
            if (!frame_.actual()) return true;
            cancel();
            int dx = e.key == Key::Left ? -1 : e.key == Key::Right ? 1 : 0, dy = e.key == Key::Up ? -1 : e.key == Key::Down ? 1 : 0;
            changed = frame_.pan(dx, dy);
            if (!changed && phase_ == Phase::Decoding) changed = true;  // it was stopped: start it again
            break;
        }
        case Key::Char:
            if (e.ch != 'i' && e.ch != 'I') return false;
            if (phase_ == Phase::Shown) say(describe());
            return true;
        default: return false;
    }
    if (changed) {
        phase_ = Phase::Wanted;
        note_.clear();
        noteDrawn_ = false;
    }
    return true;
}

// The strip the note was written over goes back as it was: no decoding for that.
void ImagePane::hideNote(Canvas* c) {
    if (noteDrawn_ && c && phase_ == Phase::Shown && under_.size() == static_cast<size_t>(c->width()) * kNoteHeight) {
        const auto& area = theme::kContent;
        uint8_t* screen = static_cast<uint8_t*>(c->getBuffer());
        std::memcpy(screen + (area.y + area.h - kNoteHeight) * c->width(), under_.data(), under_.size());
    }
    noteDrawn_ = false;
    note_.clear();
}

bool ImagePane::update(uint32_t) {
    if (path_.empty()) return false;
    uint32_t now = millis();
    if (phase_ == Phase::Wanted) return true;
    if (phase_ == Phase::Decoding) {
        if (job_ && job_->done) return true;
        if (now - pushedMs_ < kPushMs) return false;
        pushedMs_ = now;
        return true;  // what has arrived so far
    }
    if (!note_.empty() && now - noteMs_ >= kNoteMs) {
        hideNote(canvas_);
        return true;
    }
    return !note_.empty() && !noteDrawn_;
}

void ImagePane::start(Canvas& c) {
    cancel();
    job_ = std::make_shared<Job>();
    job_->path = path_;
    job_->info = info_;
    job_->frame = frame_;
    job_->map = frame_.map();
    job_->size = size_;
    job_->shotAt = shotAt_;
    job_->screen = static_cast<uint8_t*>(c.getBuffer());
    job_->stride = c.width();
    auto job = job_;
    storage_.runJob([job]() { job->run(); });
    phase_ = Phase::Decoding;
    pushedMs_ = millis();
}

void ImagePane::draw(Canvas& c) {
    if (path_.empty()) return;
    canvas_ = &c;
    const auto& area = theme::kContent;
    if (phase_ == Phase::Wanted) {
        c.fillRect(area.x, area.y, area.w, area.h, theme::kBackground);
        noteDrawn_ = false;
        problem_.clear();
        start(c);
        return;
    }
    if (phase_ == Phase::Decoding) {
        if (!job_ || !job_->done) return;  // the picture is arriving in the buffer by itself
        problem_ = job_->why;
        job_.reset();
        phase_ = Phase::Shown;
        if (!problem_.empty()) {
            c.fillRect(area.x, area.y, area.w, area.h, theme::kBackground);
            c.setFont(&fonts::body);
            c.setTextColor(theme::kMuted);
            c.setTextDatum(middle_center);
            c.drawString(problem_.c_str(), area.x + area.w / 2, area.y + area.h / 2);
            c.setTextDatum(top_left);
        } else if (!told_) {
            told_ = true;
            note_ = describe();
            noteMs_ = millis();
        }
    }
    if (!note_.empty() && !noteDrawn_) {
        int top = area.y + area.h - kNoteHeight;
        uint8_t* screen = static_cast<uint8_t*>(c.getBuffer());
        under_.assign(screen + top * c.width(), screen + (top + kNoteHeight) * c.width());
        c.fillRect(area.x, top, area.w, kNoteHeight, theme::kBackground);
        c.setFont(&fonts::small);
        c.setTextColor(theme::kText);
        c.setTextDatum(top_left);
        c.drawString(note_.c_str(), area.x + 4, top + 2);
        noteDrawn_ = true;
    }
}

}  // namespace roro