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

#include "LiveCollection.h"

#include "../writer/HDF5NXmx.h"

void LiveCollection::Expect() {
    {
        std::unique_lock ul(m);
        const uint64_t generation = state.generation + 1;
        state = Snapshot{};
        state.generation = generation;
        start_cbor.clear();
    }
    cv.notify_all();
}

void LiveCollection::Begin(const StartMessage &start, std::vector<uint8_t> cbor) {
    {
        std::unique_lock ul(m);
        // The collection Expect announced, or - when nothing announced it - a new one.
        if (state.generation == 0 || state.started || state.ended) {
            const uint64_t generation = state.generation + 1;
            state = Snapshot{};
            state.generation = generation;
        }
        state.started = true;
        state.run_number = start.run_number;
        state.run_name = start.run_name;
        state.file_prefix = start.file_prefix;
        state.number_of_images = start.number_of_images;
        state.images_per_file = start.images_per_file;
        start_message = start;
        start_cbor = std::move(cbor);
    }
    cv.notify_all();
}

void LiveCollection::FileClosed(uint64_t file_number, uint64_t image_count) {
    {
        std::unique_lock ul(m);
        if (!state.started || state.ended)
            return;
        state.files.push_back(DataFile{
            .file_number = file_number,
            .file = HDF5Metadata::DataFileName(start_message, static_cast<int64_t>(file_number)),
            .first_image = file_number * static_cast<uint64_t>(state.images_per_file),
            .image_count = image_count
        });
    }
    cv.notify_all();
}

void LiveCollection::End(uint64_t images, const std::string &error) {
    {
        std::unique_lock ul(m);
        if (state.generation == 0 || state.ended)
            return;
        state.ended = true;
        state.end_images = images;
        state.end_error = error;
    }
    cv.notify_all();
}

void LiveCollection::EndIfOpen(const std::string &error) {
    End(0, error);
}

std::optional<std::vector<uint8_t>> LiveCollection::StartCBOR() const {
    std::unique_lock ul(m);
    if (!state.started)
        return {};
    return start_cbor;
}

LiveCollection::Snapshot LiveCollection::Wait(uint64_t generation, bool started_seen, size_t files_seen,
                                              std::chrono::milliseconds wait) {
    std::unique_lock ul(m);
    cv.wait_for(ul, wait, [&] {
        return state.generation != generation || state.ended || state.files.size() > files_seen
               || (state.started && !started_seen);
    });
    return state;
}

LiveCollection::Snapshot LiveCollection::Get() const {
    std::unique_lock ul(m);
    return state;
}

bool LiveCollection::AcquireFollower() {
    std::unique_lock ul(m);
    if (follower)
        return false;
    follower = true;
    return true;
}

void LiveCollection::ReleaseFollower() {
    std::unique_lock ul(m);
    follower = false;
}