#include "aare/ZmqSocket.hpp"
#include <zmq.h>
#include <fmt/core.h>

namespace aare{


ZmqSocket::ZmqSocket(const std::string& endpoint):m_endpoint(endpoint){
    memset(m_header_buffer, 0, m_max_header_size);
}

void ZmqSocket::connect(){
    m_context = zmq_ctx_new();
    m_socket = zmq_socket(m_context, ZMQ_SUB);
    fmt::print("Setting ZMQ_RCVHWM to {}\n", m_zmq_hwm);
    int rc = zmq_setsockopt(m_socket, ZMQ_RCVHWM, &m_zmq_hwm, sizeof(m_zmq_hwm)); //should be set before connect
    if (rc)
        throw std::runtime_error(fmt::format("Could not set ZMQ_RCVHWM: {}", strerror(errno)));
    
    int bufsize = 1024*1024*m_zmq_hwm;
    fmt::print("Setting ZMQ_RCVBUF to: {} MB\n", bufsize/(1024*1024));
    rc = zmq_setsockopt(m_socket, ZMQ_RCVBUF, &bufsize, sizeof(bufsize));
    if (rc)
        throw std::runtime_error(fmt::format("Could not set ZMQ_RCVBUF: {}", strerror(errno)));

    zmq_connect(m_socket, m_endpoint.c_str());
    zmq_setsockopt(m_socket, ZMQ_SUBSCRIBE, "", 0);
}

void ZmqSocket::disconnect() {
    zmq_close(m_socket);
    zmq_ctx_destroy(m_context);
    m_socket = nullptr;
    m_context = nullptr;
}

ZmqSocket::~ZmqSocket(){
    if (m_socket)
        disconnect();
    delete[] m_header_buffer;
}

void ZmqSocket::set_zmq_hwm(int hwm){
    m_zmq_hwm = hwm;
}

void ZmqSocket::set_timeout_ms(int n){
    m_timeout_ms = n;
}   

int ZmqSocket::receive(zmqHeader &header, std::byte *data){

    //receive header
    int header_bytes_received = zmq_recv(m_socket, m_header_buffer, m_max_header_size, 0);
    m_header_buffer[header_bytes_received] = '\0'; //make sure we zero terminate
    if (header_bytes_received < 0){
        fmt::print("Error receiving header: {}\n", strerror(errno));
        return -1;
    }
    fmt::print("Bytes: {}, Header: {}\n", header_bytes_received, m_header_buffer);
    
    //decode header
    if (!decode_header(header)){
        fmt::print("Error decoding header\n");
        return -1;
    }

    //do we have a multipart message (data following header)?
    int more;
    size_t more_size = sizeof(more);
    zmq_getsockopt(m_socket, ZMQ_RCVMORE, &more, &more_size);
    if (!more) {
        return 0;   //no data following header
    }else{
        int data_bytes_received = zmq_recv(m_socket, data, 1024*1024*2, 0); //TODO! configurable size!!!!
        if (data_bytes_received == -1)
            throw std::runtime_error("Got half of a multipart msg!!!");
    }   
    return 1;
}

bool ZmqSocket::decode_header(zmqHeader &h){
    //TODO: implement
    return true;
}

} // namespace aare