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

#include "TCPStreamPusher.h"


TCPStreamPusher::TCPStreamPusher(const std::vector<std::string> &addr,
                                 std::optional<int32_t> send_buffer_size,
                                 std::optional<size_t> zerocopy_threshold,
                                 size_t send_queue_size)
        : serialization_buffer(256 * 1024 * 1024),
          serializer(serialization_buffer.data(), serialization_buffer.size()) {
    if (addr.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "No TCP writer address provided");

    for (size_t i = 0; i < addr.size(); i++) {
        socket.emplace_back(std::make_unique<TCPStreamPusherSocket>(
                addr[i], static_cast<uint32_t>(i), send_buffer_size, zerocopy_threshold, send_queue_size));
    }
}


void TCPStreamPusher::StartDataCollection(StartMessage &message) {
    if (message.images_per_file < 1)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Images per file cannot be zero or negative");
    images_per_file = message.images_per_file;
    run_number = message.run_number;
    run_name = message.run_name;

    for (size_t i = 0; i < socket.size(); i++) {
        if (!socket[i]->AcceptConnection(std::chrono::seconds(5)))
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "TCP accept timeout/failure on socket " + socket[i]->GetEndpointName());
    }

    for (size_t i = 0; i < socket.size(); i++) {
        message.socket_number = static_cast<int64_t>(i);
        if (i > 0)
            message.write_master_file = false;

        serializer.SerializeSequenceStart(message);
        socket[i]->SetRunNumber(run_number);

        if (!socket[i]->Send(serialization_buffer.data(), serializer.GetBufferSize(), TCPFrameType::START))
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Timeout/failure sending START");
    }

    for (auto &s : socket)
        s->StartWriterThread();
}

bool TCPStreamPusher::SendImage(const uint8_t *image_data, size_t image_size, int64_t image_number) {
    if (socket.empty())
        return false;

    auto socket_number = (image_number / images_per_file) % socket.size();
    if (socket[socket_number]->IsBroken())
        return false;

    return socket[socket_number]->Send(image_data, image_size, TCPFrameType::DATA, image_number);
}

void TCPStreamPusher::SendImage(ZeroCopyReturnValue &z) {
    if (socket.empty()) {
        z.release();
        return;
    }

    auto socket_number = (z.GetImageNumber() / images_per_file) % socket.size();
    if (socket[socket_number]->IsBroken()) {
        z.release();
        return;
    }

    socket[socket_number]->SendImage(z);
}

bool TCPStreamPusher::EndDataCollection(const EndMessage &message) {
    serializer.SerializeSequenceEnd(message);

    bool ret = true;
    for (auto &s : socket) {
        s->StopWriterThread();
        if (s->IsBroken())
            ret = false;
        else if (!s->Send(serialization_buffer.data(), serializer.GetBufferSize(), TCPFrameType::END))
            ret = false;
    }
    return ret;
}

std::string TCPStreamPusher::Finalize() {
    std::string ret;
    if (writer_notification_socket) {
        for (size_t i = 0; i < socket.size(); i++) {
            auto n = writer_notification_socket->Receive(run_number, run_name);
            if (!n)
                ret += "Writer " + std::to_string(i) + ": no end notification received within 1 minute from collection end";
            else if (static_cast<size_t>(n->socket_number) >= socket.size())
                ret += "Writer " + std::to_string(i) + ": mismatch in socket number";
            else if (!n->ok)
                ret += "Writer " + std::to_string(i) + ": " + n->error;
        }
    }
    return ret;
}

std::string TCPStreamPusher::GetWriterNotificationSocketAddress() const {
    if (writer_notification_socket)
        return writer_notification_socket->GetEndpointName();
    else
        return "";
}

TCPStreamPusher &TCPStreamPusher::WriterNotificationSocket(const std::string &addr) {
    writer_notification_socket = std::make_unique<ZMQWriterNotificationPuller>(addr, std::chrono::minutes(1));
    return *this;
}

std::string TCPStreamPusher::PrintSetup() const {
    std::string output = "TCPStream2Pusher: Sending images to sockets: ";
    for (const auto &s : socket)
        output += s->GetEndpointName() + " ";
    return output;
}

std::string TCPStreamPusherSocket::GetEndpointName() const {
    return endpoint;
}

void TCPStreamPusherSocket::SetRunNumber(uint64_t in_run_number) {
    run_number = in_run_number;
}

bool TCPStreamPusher::SendCalibration(const CompressedImage &message) {
    if (socket.empty())
        return false;
    serializer.SerializeCalibration(message);
    return socket[0]->Send(serialization_buffer.data(), serializer.GetBufferSize(), TCPFrameType::CALIBRATION);
}