#ifndef _UDP_HPP_
#define _UDP_HPP_

#ifndef ASIO_STANDALONE
#define ASIO_STANDALONE
#endif
#include <asio.hpp>

#include <type_traits>

#include "buffers.hpp"

using asio::ip::udp;

using udp_receiver = std::function<size_t(StaticBuffer&, size_t)>;
using udp_sender = std::function<void(StaticBuffer&, size_t)>;
using udp_cb_received = std::function<void(const udp_receiver&, const udp_sender&, size_t)>;

class UDPPeer : public std::enable_shared_from_this<UDPPeer> {
public:
    UDPPeer(asio::io_context& ctx, short port) : socket(ctx, udp::endpoint(udp::v4(), port))  {}
    UDPPeer(asio::io_context& ctx) : socket(ctx, udp::v4())  {}

    void StartReceive() {
        if(started)
            return;
        started = true;
        _ReceiveFrom();
    }

    void StopReceive() {
        if(!started)
            return;
        started = false;
        socket.cancel();
    }

    void SendTo(const char* addr, uint16_t port, StaticBuffer& buf, size_t length) {
        udp::endpoint remote_peer(asio::ip::make_address_v4(addr), port);
        auto abuf = asio::buffer(buf.Data(), length);
        socket.async_send_to(abuf, remote_peer, [obuf = std::move(buf)](std::error_code ec, size_t){});
    }

    UDPPeer& SetReceivedCallback(udp_cb_received cb) { callback = cb; return *this; }

protected:
    void _ReceiveFrom() {
        socket.async_wait(udp::socket::wait_read, [this, self = this->shared_from_this()](std::error_code ec){
            if(!ec) {
                size_t reserve_size = socket.available();
                if(callback) {
                    auto receiver = [this](StaticBuffer& buf, size_t length) -> size_t {
                        return socket.receive_from(asio::buffer(buf.Data(), length), remote);
                    };
                    auto sender = [this](StaticBuffer& buf, size_t length) {
                        auto abuf = asio::buffer(buf.Data(), length);
                        socket.async_send_to(abuf, remote, [obuf = std::move(buf)](std::error_code ec, size_t){});
                    };
                    callback(receiver, sender, reserve_size);
                }
            }
            if(started)
                _ReceiveFrom();
        });
    }

    bool started = false;    
    udp_cb_received callback;
    udp::socket socket;
    udp::endpoint remote;
};

#endif