#ifndef TRANSFER_H
#define TRANSFER_H

#include <chrono>
#include <opencv2/core/mat.hpp>
#include <opencv2/imgcodecs.hpp>
#include <sys/socket.h>
#include <logging.h>
#include <vector>

using namespace std;

struct imageHeader {
    size_t size;
    chrono::milliseconds timestamp;
};

chrono::milliseconds getMillis() {
    // Get current time
    const auto now = chrono::system_clock::now();

    // Get time since epoch in milliseconds
    const chrono::duration ms = chrono::duration_cast<chrono::milliseconds>(now.time_since_epoch());

    return chrono::milliseconds(ms.count());
}

inline void serializeImage(const cv::Mat& image, std::vector<uchar>& buffer) {
    cv::imencode(".jpg", image, buffer);
}

int sendImage(int socket, const cv::Mat& image, std::vector<uchar>& buffer) {
    serializeImage(image, buffer);
    size_t totalSent = 0;

    // first send the size of the serialized image
    const size_t size = buffer.size();
    const chrono::milliseconds timestamp = getMillis();
    imageHeader header;
    header.size = size;
    header.timestamp = timestamp;
    trace("Buffer size: ", size);
    if (const ssize_t sent = send(socket, &header, sizeof(header), 0); sent == -1) {
        error("Error sending data header");
        return -1;
    }

    // then start sending the serialized image
    while (totalSent < size) {
        const ssize_t sent = send(socket, buffer.data() + totalSent, size - totalSent, 0);
        if (sent == -1) {
            error("Error sending data");
            return -1;
        }
        totalSent += sent;
        debug("Packet sent (", sent, " bytes, total ", totalSent, " bytes)");
    }
    return 0;
}

int recvImage(int socket, std::vector<uchar>& buffer) {
    // first receive the size of the image
    imageHeader header;
    if (ssize_t recvd = recv(socket, &header, sizeof(imageHeader), 0); recvd <= 0) {
        error("Error receiving data header");
        return -1;
    }

    size_t dataSize = header.size;
    chrono::milliseconds sentTime = header.timestamp;

    trace("Buffer size: ", dataSize);

    // resize the buffer to fit the whole image
    buffer.resize(dataSize);

    // start receiving the image until the whole thing arrives
    size_t totalReceived = 0;
    while (totalReceived < dataSize) {
        ssize_t bytesReceived = recv(socket, buffer.data() + totalReceived, dataSize, 0);
        if (bytesReceived <= 0) {
            error("Error receiving data");
            buffer.clear();
            return -1;
        }
        totalReceived += bytesReceived;
        debug("Packet received (", bytesReceived, " bytes, total ", totalReceived, " bytes)");
    }
    chrono::milliseconds currentTime = getMillis();

    chrono::milliseconds diff = currentTime - sentTime;
    debug("Packet latency: ", diff.count(), "ms");
    return 0;
}

bool applyImage(cv::Mat& image, std::vector<uchar> *src) {
    // decode the image into an OpenCV Mat
    cv::imdecode(*src, cv::IMREAD_UNCHANGED, &image);
    return true;
}
#endif //TRANSFER_H