// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "JFJochImageBufferReader.h"

#include <algorithm>
#include <cstring>

#include "JFJochStartMessageReader.h"
#include "../common/JFJochException.h"
#include "../frame_serialize/CBORStream2Deserializer.h"

JFJochImageBufferReader::JFJochImageBufferReader(ImageBuffer &buffer) : buffer_(buffer) {}

void JFJochImageBufferReader::Open(const StartMessage &start) {
    {
        std::unique_lock ul(m_);
        number_of_images_ = start.number_of_images;
        present_.assign(number_of_images_, false);
        max_present_ = -1;
        ended_ = false;
    }
    SetStartMessage(DatasetFromStartMessage(start, default_experiment));
}

void JFJochImageBufferReader::Ended() {
    {
        std::unique_lock ul(m_);
        ended_ = true;
    }
    cv_.notify_all();
}

uint64_t JFJochImageBufferReader::GetNumberOfImages() const {
    return number_of_images_;
}

void JFJochImageBufferReader::Close() {
    SetStartMessage({});
}

void JFJochImageBufferReader::Poll(std::unique_lock<std::mutex> &ul) {
    polling_ = true;
    next_poll_ = std::chrono::steady_clock::now() + POLL_PERIOD;
    ul.unlock();
    const std::vector<int64_t> images = buffer_.GetStatus().images_in_the_buffer;
    ul.lock();
    polling_ = false;
    for (const int64_t i: images)
        if (i >= 0 && static_cast<uint64_t>(i) < number_of_images_) {
            present_[i] = true;
            max_present_ = std::max(max_present_, i);
        }
    cv_.notify_all();
}

bool JFJochImageBufferReader::WaitForMessage(int64_t image_number, std::vector<uint8_t> &cbor) {
    if (image_number < 0 || static_cast<uint64_t>(image_number) >= number_of_images_)
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds, "Image out of bounds");
    std::unique_lock ul(m_);
    // Images arrive in order, give or take a few: one still missing when the buffer already holds an
    // image MAX_IMAGE_LAG past it is lost, not late. After the end, one is asked for once more, as it
    // may have come in just before.
    while (!present_[image_number] && !ended_ && max_present_ < image_number + MAX_IMAGE_LAG) {
        const auto now = std::chrono::steady_clock::now();
        if (!polling_ && now >= next_poll_)
            Poll(ul);
        else
            cv_.wait_until(ul, polling_ ? now + POLL_PERIOD : next_poll_);
    }
    ul.unlock();
    return buffer_.GetImage(cbor, image_number);
}

bool JFJochImageBufferReader::ReadRawImage(int64_t image_number, JFJochReaderRawImage &ret) {
    if (!WaitForMessage(image_number, ret.read_buffer))
        return false;
    const auto msg = CBORStream2Deserialize(ret.read_buffer);
    if (msg->msg_type != CBORImageType::IMAGE || !msg->data_message)
        return false;
    // Copied out of the message: read_buffer is scratch a caller may hand to the next read.
    const CompressedImage &image = msg->data_message->image;
    ret.image_buffer.resize(image.GetCompressedSize());
    memcpy(ret.image_buffer.data(), image.GetCompressed(), image.GetCompressedSize());
    ret.image = CompressedImage(ret.image_buffer.data(), ret.image_buffer.size(), image.GetWidth(),
                                image.GetHeight(), image.GetMode(), image.GetCompressionAlgorithm());
    return true;
}

bool JFJochImageBufferReader::LoadImage_i(std::shared_ptr<JFJochReaderDataset> &dataset, DataMessage &message,
                                          std::vector<uint8_t> &buffer, int64_t image_number, bool) {
    if (!dataset || !WaitForMessage(image_number, buffer))
        return false;
    const auto msg = CBORStream2Deserialize(buffer);
    if (msg->msg_type != CBORImageType::IMAGE || !msg->data_message)
        return false;
    message = *msg->data_message;
    return true;
}