#include <iostream>

#include <asio.hpp>
#include <blt/std/logging.h>
#include <blt/std/memory.h>
#include <string>
#include <array>
#include <iomanip> // This might be necessary

using asio::ip::udp;

class byte_reader
{
    private:
        unsigned char* data;
        size_t current_byte = 0;
    public:
        explicit byte_reader(unsigned char* data): data(data)
        {
            reset();
        }
        
        inline void reset()
        {
            // magic number for end of header
            current_byte = 12;
        }
        
        inline unsigned char next()
        {
            return data[current_byte++];
        }
};

int main()
{
    try
    {
        asio::io_context io_context;
        
        udp::socket socket(io_context, udp::endpoint(udp::v4(), 5555));
        
        while (true)
        {
            std::array<unsigned char, 65535> recv_buf{};
            udp::endpoint remote_endpoint;
            size_t bytes = socket.receive_from(asio::buffer(recv_buf), remote_endpoint);
            
            // UDP data contains:
            // 2 bytes for transaction id
            // 2 bytes for flags
            // 2 bytes for number of questions
            // 2 bytes for Answer RRs
            // 2 bytes for Authority RRs
            // 2 bytes for Additional RRs
            
            // question format:
            // 1 byte for length
            //
            
            // get the number of questions
            uint16_t questions; // yes I made this part of my library just for this :3
            blt::mem::fromBytes(&recv_buf[4], questions); // i hate little endian
            
            byte_reader reader(recv_buf.data());
            
            BLT_INFO("Bytes received: %d with %d questions", bytes, questions);
            
            for (uint16_t i = 0; i < questions; i++)
            {
                std::string domain;
                while (true)
                {
                    uint8_t length = reader.next();
                    if (length == 0)
                        break;
                    if (!domain.empty())
                        domain += '.';
                    for (uint8_t j = 0; j < length; j++)
                        domain += static_cast<char>(reader.next());
                }
                BLT_INFO("DOMAIN: %s", domain.c_str());
            }

//            for (size_t i = 0; i < bytes; i++)
//            {
//                std::cout << recv_buf[i] << " ";
//            }
//            std::cout << std::endl;
//            for (size_t i = 0; i < bytes; i++)
//            {
//                std::cout << std::hex << std::uppercase << ((unsigned short)recv_buf[i]) << std::nouppercase << std::dec << " ";
//            }
//            std::cout << std::endl;
            
            std::string message = "Hello!";
            
            asio::error_code ignored_error;
            socket.send_to(asio::buffer(message),
                           remote_endpoint, 0, ignored_error);
        }
    }
    catch (std::exception& e)
    {
        std::cerr << e.what() << std::endl;
    }
    
    std::cout << "Hello, World!" << std::endl;
    return 0;
}
