/**
 * Brett    Terpstra    6920201
 * Michael  Boulos      6973523
 * Adrian   Pinu        6970677
 */

#include <iostream>

#include <asio.hpp>
#include <blt/std/logging.h>
#include <blt/std/memory.h>
#include <string>
#include <array>
#include <thread>
#include <iomanip> // This might be necessary
#include <vector>
#include <type_traits>
#include <unordered_set>
#include "ip.h"

/**
 * ----------------------------
 *          | CONFIG |
 * ----------------------------
 */
// should we strictly match results? ie block *wikipedia.org* or just wikipedia.org?
static constexpr bool STRICT_MATCHING = false;
// true -> only match A records ; false -> match any named record (configure with NON_STRICT_REPLACE_ALL)
static constexpr bool STRICT_FILTERING = false;
// true -> match all records ; false -> match only records we might want to replace (A, AAAA, CNAME)
static constexpr bool NON_STRICT_REPLACE_ALL = true;

// DNS server to use for forwarding to / resolving DNS requests
static inline constexpr std::string DNS_SERVER_IP()
{
    return "8.8.8.8";
}

// replacement IP address. Make sure this is a 4 octet string seperated by .
static inline constexpr IPAddress REPLACEMENT_IP()
{
    return {"139.57.100.6"};
}

static const std::unordered_set<std::string> DISALLOWED_DOMAINS{
        "en.wikipedia.org",
        "tpgc.me",
        "zombo.com"
};

// DNS 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
// (length) bytes per section
// ... until length 0
// 2 bytes for QTYPE (useless)
// 2 bytes for QCLASS (useless)

// answer format:
// 2 byte for domain name (offset ptr, still not 100 on from where)
// 2 byte for type
// 2 byte for class
// 4 byte for time to live
// 2 byte for length of data
// (lengthy) byte for data

class byte_reader
{
    private:
        unsigned char* data;
        mutable size_t current_byte = 0;
    public:
        explicit byte_reader(unsigned char* data): data(data)
        {
            reset();
        }
        
        inline void reset() const
        {
            // magic number for end of header
            current_byte = 12;
        }
        
        inline unsigned char& next() const
        {
            return data[current_byte++];
        }
        
        template<typename T>
        inline void to(T& t) const
        {
            blt::mem::fromBytes(&data[current_byte], t);
            skip(sizeof(T));
        }
        
        inline void copy(unsigned char*& out, size_t size) const
        {
            std::memcpy(out, &data[current_byte], size);
            current_byte += size;
        }
        
        inline void skip(size_t s = 1) const
        {
            current_byte += s;
        }
        
        inline size_t last() const
        {
            return current_byte;
        }
        
        inline unsigned char* from()
        {
            return &data[current_byte];
        }
};

class send_buffer;

class question
{
        friend send_buffer;
    private:
        std::string domain;
        uint16_t QTYPE;
        uint16_t QCLASS;
    public:
        explicit question(const byte_reader& reader)
        {
            // process the full question
            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());
            }
            // Skip QTYPE and QCLASS
            // but keep a copy of it for future writing
            reader.to(QTYPE);
            reader.to(QCLASS);
        }
        
        const std::string& operator()()
        {
            return domain;
        }
};

class answer
{
        friend send_buffer;
    private:
        mutable uint16_t NAME = 0;
        uint16_t TYPE = 0;
        uint16_t CLASS = 0;
        uint32_t TTL = 0;
        uint16_t RDLENGTH = 0;
        bool requires_reset = false;
        unsigned char* RDATA = nullptr;
    public:
        explicit answer(const byte_reader& reader)
        {
            reader.to(NAME);
            reader.to(TYPE);
            reader.to(CLASS);
            reader.to(TTL);
            reader.to(RDLENGTH);
            RDATA = new unsigned char[RDLENGTH];
            reader.copy(RDATA, RDLENGTH);
            
            BLT_TRACE("%d, %d, %d, %d, %d", NAME, TYPE, CLASS, TTL, RDLENGTH);
        }
        
        [[nodiscard]] uint16_t type() const
        {
            return TYPE;
        }
        
        inline void substitute(const IPAddress& addr)
        {
            BLT_DEBUG("Substituting with replacement address '%s'", REPLACEMENT_IP().asString.c_str());
            BLT_ASSERT(RDLENGTH == 4);
            std::memcpy(RDATA, addr.octets, 4);
        }
        
        inline void setARecord(const IPAddress& addr)
        {
            BLT_DEBUG("Setting to A record");
            NAME = 0;
            NAME |= (0b11 << 14);
            requires_reset = true;
            BLT_INFO(NAME);
            delete[] RDATA;
            RDATA = new unsigned char[4];
            RDLENGTH = 4;
            TYPE = 1;
            CLASS = 1;
            substitute(addr);
        }
        
        inline void reset(size_t offset) const
        {
            auto i16 = static_cast<uint16_t>(offset) & (~(0b11 << 14));
            NAME |= i16;
        }
        
        // rule of 5
        answer(const answer& answer) = delete;
        
        answer& operator=(const answer& answer) = delete;
        
        answer(answer&& move)
        {
            NAME = move.NAME;
            TYPE = move.TYPE;
            CLASS = move.CLASS;
            TTL = move.TTL;
            RDLENGTH = move.RDLENGTH;
            RDATA = move.RDATA;
            move.RDATA = nullptr;
        }
        
        answer& operator=(answer&& move)
        {
            NAME = 0;
            NAME = move.NAME;
            TYPE = move.TYPE;
            CLASS = move.CLASS;
            TTL = move.TTL;
            RDLENGTH = move.RDLENGTH;
            RDATA = move.RDATA;
            move.RDATA = nullptr;
            return *this;
        }
        
        ~answer()
        {
            delete[] RDATA;
        }
};

class send_buffer
{
    private:
        mutable std::array<unsigned char, 65535> internal_data{};
        mutable size_t write_index = 0;
    public:
        send_buffer() = default;
        
        void write(unsigned char* data, size_t size) const
        {
            std::memcpy(&internal_data[write_index], data, size);
            write_index += size;
        }
        
        template<typename T>
        void write(const T& t) const
        {
            if constexpr (std::is_arithmetic_v<T>)
            {
                blt::mem::toBytes(t, &internal_data[write_index]);
                write_index += sizeof(T);
            } else if constexpr (std::is_same_v<T, std::string>)
            {
                static_assert("No");
            } else if constexpr (std::is_same_v<T, answer>)
            {
                write(t.NAME);
                write(t.TYPE);
                write(t.CLASS);
                write(t.TTL);
                write(t.RDLENGTH);
                std::memcpy(&internal_data[write_index], t.RDATA, t.RDLENGTH);
                write_index += t.RDLENGTH;
            } else if constexpr (std::is_same_v<T, question>)
            {
                auto labels = blt::string::split(t.domain, '.');
                for (const auto& label : labels)
                {
                    auto length = static_cast<uint8_t>(label.length());
                    write(length);
                    std::memcpy(&internal_data[write_index], &label[0], length);
                    write_index += length;
                }
                // write length of 0 to signal end of labels
                write('\0');
                write(t.QTYPE);
                write(t.QCLASS);
            } else
                static_assert("Data type not supported!");
        }
        
        void reset() const
        {
            write_index = 0;
        }
        
        [[nodiscard]] unsigned char operator[](size_t i)
        {
            return internal_data[i];
        }
        
        [[nodiscard]] unsigned char* data()
        {
            return internal_data.data();
        }
        
        [[nodiscard]] size_t size() const
        {
            return write_index;
        }
        
        auto buffer() const
        {
            return asio::buffer(internal_data.data(), size());
        }
};

inline bool shouldReplace(const answer& a)
{
    // a records will be handled in either case, check for others like AAAA or CNAME
    // TODO: add enums to this + a way to add custom types
    return NON_STRICT_REPLACE_ALL || a.type() == 28 || a.type() == 5;
}

void process_answers(std::vector<answer>& answers)
{
    for (auto& a : answers)
    {
        if (a.type() == 1)
        {
            a.substitute(REPLACEMENT_IP());
        } else if (!STRICT_FILTERING && shouldReplace(a))
        {
            a.setARecord(REPLACEMENT_IP());
        }
    }
}

int main()
{
    try
    {
        asio::io_context io_context;
        
        udp::socket socket(io_context, udp::endpoint(udp::v6(), 5555));
        
        std::array<unsigned char, 65535> recv_buf{};
        std::array<unsigned char, 65535> mod_recv_buf{};
        while (true)
        {
            udp::endpoint remote_endpoint;
            size_t bytes = socket.receive_from(asio::buffer(recv_buf), remote_endpoint);
            
            // 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
            
            BLT_INFO("Bytes received: %d with %d questions", bytes, questions);
            
            // forward to google.
            size_t out_bytes;
            sendUDPMessage(DNS_SERVER_IP(), recv_buf.data(), bytes, mod_recv_buf, out_bytes);
            
            uint16_t num_of_answers;
            blt::mem::fromBytes(&mod_recv_buf[6], num_of_answers);
            
            byte_reader reader2(mod_recv_buf.data());
            
            BLT_INFO("Bytes answered %d with %d answers", out_bytes, num_of_answers);
            
            // no one actually does multiple questions. trying to do it in dig is not easy
            // and the standard isn't really designed for this (how do we handle if one question errors but the other doesn't? there is only
            // one return code.)
            question q(reader2);
            std::vector<answer> answers;
            for (int i = 0; i < num_of_answers; i++)
            {
                answer a(reader2);
                answers.push_back(std::move(a));
            }
            
            BLT_INFO("DOMAIN: %s", q().c_str());
            if (STRICT_MATCHING && DISALLOWED_DOMAINS.contains(q()))
                process_answers(answers);
            else if (!STRICT_MATCHING)
            {
                // linear search the domains for contains. Maybe find a better way to do this.
                for (const auto& v : DISALLOWED_DOMAINS)
                    if (blt::string::contains(q(), v))
                        process_answers(answers);
            }
            
            send_buffer send;
            send.write(mod_recv_buf.data(), 12);
            auto question_offset = send.size();
            send.write(q);
            for (const answer& a : answers)
            {
                BLT_TRACE("Writing answer with type of %d", a.type());
                a.reset(question_offset);
                send.write(a);
            }
            send.write(reader2.from(), out_bytes - reader2.last());
            
            asio::error_code ignored_error;
            socket.send_to(send.buffer(), remote_endpoint, 0, ignored_error);
        }
    }
    catch (std::exception& e)
    {
        std::cerr << e.what() << std::endl;
    }
    
    std::cout << "Hello, World!" << std::endl;
    return 0;
}
