#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>
#include <boost/asio.hpp>

using namespace std;
using boost::asio::ip::tcp;

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);
}

void sendImage(tcp::socket& socket, const cv::Mat& image, std::vector<uchar>& buffer) {
    boost::system::error_code e;

    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 = boost::asio::write(socket, boost::asio::buffer(&header, sizeof(header)), e); sent == -1) {
        throw boost::system::system_error(e);
    }

    // then start sending the serialized image
    while (totalSent < size) {
        const ssize_t sent = boost::asio::write(socket, boost::asio::buffer(buffer), e);
        if (sent == -1) {
            throw boost::system::system_error(e);
        }
        totalSent += sent;
        debug("Packet sent (", sent, " bytes, total ", totalSent, " / ", size, " bytes)");
    }
}

int recvImage(tcp::socket& socket, std::vector<uchar>& buffer) {
    boost::system::error_code e;

    // first receive the size of the image
    imageHeader header;

    size_t len = 0;
    while (len < sizeof(imageHeader)) {
        len += socket.read_some(boost::asio::buffer(&header, sizeof(header)), e);
        if (e.failed()) {
            throw boost::system::system_error(e);
        }
    }

    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) {
        size_t bytesReceived = boost::asio::read(socket, boost::asio::buffer(buffer), e);
        if (e == boost::asio::error::eof) // Connection closed cleanly by peer.
            throw boost::system::system_error(e);
        else if (e)
            throw boost::system::system_error(e);
        totalReceived += bytesReceived;
        debug("Packet received (", bytesReceived, " bytes, total ", totalReceived, " / ", dataSize, " bytes)");
    }
    chrono::milliseconds currentTime = getMillis();

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

    return diff.count();
}

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