#include <iostream>
#include <unordered_map>
#include <thread>
#include <tins/tins.h>
#include <thread>
#include <iterator>
#include <netinet/in.h>
#include <linux/if_packet.h>
#include <sys/stat.h>
#include <fstream>
#include <utility>
#include <regex>
#include <atomic>
#include "log4cpp/Category.hh"
#include "log4cpp/Appender.hh"
#include "log4cpp/OstreamAppender.hh"
#include "log4cpp/Layout.hh"
#include "log4cpp/BasicLayout.hh"
#include "log4cpp/Priority.hh"
#include "log4cpp/PatternLayout.hh"
#include <net/if.h>
#include <net/ethernet.h>
#include <stdio.h>
#include <sys/socket.h>
#include <unistd.h>

using namespace std;

log4cpp::Category& logroot = log4cpp::Category::getRoot();

struct iface_addr {
    Tins::IPv4Address ipaddr;
    Tins::IPv4Address netmask;
    Tins::IPv4Address gateway;
    bool multiple_addrs = false;
};

typedef std::unordered_map<int, iface_addr> iface_addrs;

struct _addr_table {
    std::array<iface_addrs, 2> ipaddrs;
    std::array<iface_addrs, 2> stbaddrs;
    std::atomic<int> active_index;
} addr_table;

struct network {
    Tins::IPv4Range adrange = Tins::IPv4Range(0,0); // Subscriber ipaddress
    Tins::IPv4Address netmask; // Netmask
    Tins::IPv4Address gateway;
    bool stb = false;
};

struct stboption {
    std::string vendor_class;
    std::vector<uint8_t> opt43;
};

struct _config { // Global struct holding the config
    // MAIN
    std::vector<network> networks; // Supernets
    std::vector<Tins::IPv4Address> nameservers;

    uint32_t lease_time;

    // STB DHCP Vendor Options
    std::vector<stboption> stboptions;
} config;

struct request {
    int ifindex = 0;
    bool broadcast = false;
    uint32_t xid = 0;
    Tins::IPv4Address req_ip = 0;
    Tins::IPv4Address ciaddr = 0;
    Tins::HWAddress<ETHER_ADDR_LEN> chaddr;
    uint8_t message_type = 0;
    uint16_t secs;
    std::string vendor_class;
};

void maintain_addr_map() {
    int new_index = addr_table.active_index ^ 0x1;
    FILE * routefile;
    while(true) {
        iface_addrs new_ipaddrs, new_stbaddrs;
        logroot.debug("Refreshing address map");
        char buf[256];
        char ifname[IF_NAMESIZE];
        uint32_t raw_ipaddr, prefix;
        routefile = fopen("/proc/net/route","r");
        fgets(buf, sizeof(buf), routefile);
        while (fgets(buf, sizeof(buf), routefile) != NULL) {
            if (sscanf(buf, "%s\t%x\t%*x\t%*x\t%*x\t%*x\t%*x\t%x", ifname, &raw_ipaddr, &prefix) == 3) {
                if (prefix != 0xFFFFFFFF) continue; // Not a host route. Skip.
                int ifindex = if_nametoindex(ifname);
                if (!ifindex) {
                    logroot.error("Iface index not found: %s", ifname);
                    continue;
                }

                iface_addr new_addr;
                Tins::IPv4Address ipaddr(raw_ipaddr);
                for( network net : config.networks) {
                    if (net.adrange.contains(ipaddr)) {
                        new_addr.ipaddr = ipaddr;
                        new_addr.netmask = net.netmask;
                        new_addr.gateway = net.gateway;
                        iface_addrs &addrs = (!net.stb) ? new_ipaddrs : new_stbaddrs;
                        iface_addrs::const_iterator iter = addrs.find(ifindex);
                        if (iter == addrs.end()) {
                            addrs[ifindex] = new_addr;
                        } else {
                            addrs[ifindex].multiple_addrs = true;
                        }

                        break;
                    }
                }
            }
        }
        fclose(routefile);
        addr_table.ipaddrs[new_index].swap(new_ipaddrs);
        addr_table.stbaddrs[new_index].swap(new_stbaddrs);
        addr_table.active_index = new_index;
        logroot.debug("Addrs: %i",addr_table.ipaddrs[addr_table.active_index].size());
        std::this_thread::sleep_for(std::chrono::seconds(60));
    }
}

int parse_config() {
    std::vector<std::regex> cfg_regexes = {
        std::regex(R"(^(lease_time)=(\d+))"),
        std::regex(R"(^(nameserver)=(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}))"),
        std::regex(R"(^(network)=(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})\/)"
        R"((\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})\/)"
        R"((\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}))"),
        std::regex(R"(^(stb_network)=(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})\/)"
        R"((\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})\/)"
        R"((\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}))"),
        std::regex(R"(^(vendor)=\"([^\"]+)\":([0-9A-Fa-f]+))"),
    };
    std::smatch m;
    struct stat statbuf;
    std::string cfgfile = "/etc/nasdhcp.conf";
    std::ifstream cfg;
    if (stat(cfgfile.c_str(), &statbuf) != 0) {
        cfgfile = "vpcdhcp.conf";
        if (stat(cfgfile.c_str(), &statbuf) != 0) {
            logroot.fatal("No config file found!");
            return -1; /* No config file found */
        }
    }

    cfg.open(cfgfile, std::ios::in);

    std::string line;
    while (std::getline(cfg, line)) {
        for (std::regex re : cfg_regexes) {
            std::regex_search(line, m, re);
            if (m.size() > 1) {
                if (m[1] == "lease_time") {
                    config.lease_time = (uint32_t) std::stoul(m[2]);
                }
                else if (m[1] == "nameserver") {
                    config.nameservers.push_back(Tins::IPv4Address(m[2]));
                }
                else if (m[1] == "network") {
                    network net;
                    net.adrange = Tins::IPv4Range::from_mask(Tins::IPv4Address(m[2]), Tins::IPv4Address(m[3]));
                    net.netmask = Tins::IPv4Address(m[3]);
                    net.gateway = Tins::IPv4Address(m[4]);
                    config.networks.push_back(net);
                }
                else if (m[1] == "stb_network") {
                    network net;
                    net.adrange = Tins::IPv4Range::from_mask(Tins::IPv4Address(m[2]), Tins::IPv4Address(m[3]));
                    net.netmask = Tins::IPv4Address(m[3]);
                    net.gateway = Tins::IPv4Address(m[4]);
                    net.stb = true;
                    config.networks.push_back(net);
                }
                else if (m[1] == "vendor") {
                    if (m[3].length() % 2) {
                        logroot.error("Vendor data: number of bytes must be even");
                        return -1;
                    }
                    stboption sopt;
                    sopt.vendor_class = m[2];
                    char hexbyte[2];
                    uint8_t temp;
                    std::stringstream ss(m[3]);

                    while (ss.read(hexbyte,2)) {
                        sscanf(hexbyte,"%02hhx", &temp);
                        sopt.opt43.push_back(temp);
                    }
                    config.stboptions.push_back(sopt);
                }
                break;
            }
        }
    }
    cfg.close();
    return 0;
}

std::string msgtype(int i) {
    switch (i) {
        case Tins::DHCP::DISCOVER:
            return "DISCOVER";
            break;
        case Tins::DHCP::REQUEST:
            return "REQUEST";
            break;
        case Tins::DHCP::OFFER:
            return "OFFER";
            break;
        case Tins::DHCP::ACK:
            return "ACK";
            break;
        case Tins::DHCP::NAK:
            return "NAK";
            break;
        default:
            return "UNKNOWN";
            break;
    }
}

int recv_dhcp_packet(int sockrcv, request &req) {
    uint8_t buf[65536];
    uint8_t aux_buf[CMSG_SPACE(2048)];
    struct iovec riov[1];
    riov[0].iov_base = buf;
    riov[0].iov_len = sizeof(buf);

    struct msghdr msg;
    memset((void *) &msg, 0, sizeof(msg));
    msg.msg_iov = riov;
    msg.msg_iovlen = 1;
    msg.msg_control = aux_buf;
    msg.msg_controllen = sizeof(aux_buf);

    int bytesrcv = recvmsg(sockrcv, &msg, 0);
    if (bytesrcv > 0) {
        // Retrieve ifindex from ancillary data
        struct cmsghdr *cmsg;
        for (cmsg = CMSG_FIRSTHDR(&msg); cmsg != NULL; cmsg = CMSG_NXTHDR(&msg, cmsg)) {
            if ((cmsg->cmsg_level == IPPROTO_IP) && (cmsg->cmsg_type == IP_PKTINFO)) {
                struct in_pktinfo *pinfo = (struct in_pktinfo *)CMSG_DATA(cmsg);
                req.ifindex = pinfo->ipi_ifindex;
//                req.broadcast = (pinfo->ipi_addr.s_addr == INADDR_BROADCAST);
            }
        }

        // Process DHCP Request
        Tins::DHCP dhcp_packet;
        try {
            dhcp_packet = Tins::DHCP((uint8_t *)riov[0].iov_base, bytesrcv);
        }
        catch (Tins::malformed_packet) {
            logroot.error("Malformed packet received");
            return -1;
        }
        catch (...) {
            return -1; // DHCP server should not crash because of some garbage data sent to it
        }
        if (dhcp_packet.opcode() != Tins::BootP::BOOTREQUEST) {
            return -1; // Packet is not request
        }
        // Get relevant field from request
        try {
            req.message_type = dhcp_packet.type();
        }
        catch (Tins::option_not_found) {
            return -1; // Non-DHCP packet
        }
        if ((req.message_type != Tins::DHCP::DISCOVER) && (req.message_type != Tins::DHCP::REQUEST)) {
            return -1; // Server should only reply to DISCOVER and REQUEST packets
        }

        req.xid = dhcp_packet.xid();
        req.ciaddr = dhcp_packet.ciaddr();
        req.secs = dhcp_packet.secs();
        req.broadcast = (dhcp_packet.padding() == 0x8000);
        try {
            req.req_ip = dhcp_packet.requested_ip();
        }
        catch (Tins::option_not_found) {
            req.req_ip = 0;
        }
        req.chaddr = dhcp_packet.chaddr();
        const Tins::DHCP::option *opt60 = dhcp_packet.search_option(Tins::DHCP::VENDOR_CLASS_IDENTIFIER);
        if (opt60 && (opt60->data_size() > 0)) {
            req.vendor_class = std::string((char *)(opt60->data_ptr()), opt60->data_size());
        }
    }
    logroot.infoStream() << "Received " << msgtype(req.message_type) << " on iface " << req.ifindex
                         << ": xid: " << std::hex << req.xid << " chaddr: " << req.chaddr.to_string()
                         << " req-ip: " << req.req_ip.to_string()
                         << " ciaddr: " << req.ciaddr.to_string()
                         << " vendor-class: " << req.vendor_class;
    return 0;
}

const Tins::DHCP::option opt0(0);

int dhcp_send_reply(int sockraw, request &req) {
    std::vector<stboption>::const_iterator opt43it;
    bool is_stb = false;
    for (opt43it = config.stboptions.cbegin(); opt43it != config.stboptions.cend(); ++opt43it) {
        if (opt43it->vendor_class == req.vendor_class) {
            is_stb = true;
            break;
        }
    }

    const iface_addrs &addr_map = (!is_stb) ? addr_table.ipaddrs[addr_table.active_index] :
            addr_table.stbaddrs[addr_table.active_index];

    iface_addrs::const_iterator iter = addr_map.find(req.ifindex);
    if (iter == addr_map.end()) {
        logroot.debug("No address found for iface %i", req.ifindex);
        return -1; // No ip addresses found for this iface. Ignore packet.
    }
    iface_addr iaddr = iter->second;
    if(iaddr.multiple_addrs) { // Multiple ip addresses on iface. Ignore packet.
        logroot.debug("Multiple address found for iface %i", req.ifindex);
        return -1;
    }
    Tins::IPv4Address &yiaddr = iaddr.ipaddr;
    Tins::IPv4Address &siaddr = iaddr.gateway;

    auto pkt = Tins::IP(iaddr.ipaddr, siaddr) / Tins::UDP(68, 67) / Tins::DHCP();

    if (req.broadcast)
        pkt.dst_addr(Tins::IPv4Address::broadcast);

    // DHCP Header
    Tins::DHCP &dhcpdata = pkt.rfind_pdu<Tins::DHCP>();
    dhcpdata.opcode(Tins::BootP::BOOTREPLY);
    dhcpdata.xid(req.xid);
    dhcpdata.secs(req.secs);
    dhcpdata.chaddr(req.chaddr);

    // DHCP Options
    if(req.message_type == Tins::DHCP::DISCOVER) {
        dhcpdata.type(Tins::DHCP::OFFER);
    }
    else if(req.message_type == Tins::DHCP::REQUEST) {
        if ((req.req_ip && (req.req_ip != iaddr.ipaddr)) || (req.ciaddr && (req.ciaddr != iaddr.ipaddr))) {
            dhcpdata.type(Tins::DHCP::NAK);
            pkt.dst_addr(Tins::IPv4Address::broadcast);
        }
        else {
            dhcpdata.type(Tins::DHCP::ACK);
        }
    }

    if (dhcpdata.type() != Tins::DHCP::NAK) {
        dhcpdata.yiaddr(yiaddr);
        dhcpdata.siaddr(siaddr);
        dhcpdata.subnet_mask(iaddr.netmask);
        dhcpdata.routers({siaddr});
        dhcpdata.domain_name_servers(config.nameservers);
        dhcpdata.lease_time(config.lease_time);

        if(is_stb) {
            Tins::DHCP::option opt43(Tins::DHCP::VENDOR_ENCAPSULATED_OPTIONS, opt43it->opt43.size(), opt43it->opt43.data());
            dhcpdata.add_option(opt43);
        }

    }
    dhcpdata.server_identifier(Tins::IPv4Address(siaddr));
    dhcpdata.end();

    // PAD DHCP packet to at least 300 bytes
    int pad = 300 - dhcpdata.size();
    while (pad > 0) {
        dhcpdata.add_option(opt0);
        pad -= 2;
    }

    Tins::byte_array buf = pkt.serialize();

    // Send reply packet
    struct sockaddr_ll rawaddr;
    memset((void *)&rawaddr,0,sizeof(rawaddr));
    rawaddr.sll_ifindex = req.ifindex;
    rawaddr.sll_family = AF_PACKET;
    rawaddr.sll_protocol = htons(ETH_P_IP);
    rawaddr.sll_halen = ETHER_ADDR_LEN;
    if (pkt.dst_addr() == Tins::IPv4Address::broadcast) {
        memset(rawaddr.sll_addr, 0xFF, ETHER_ADDR_LEN);
    }
    else {
        req.chaddr.copy(std::begin(rawaddr.sll_addr));
    }

    int bytes_sent = sendto(sockraw, buf.data(), buf.size(), 0, (struct sockaddr *) &rawaddr, sizeof(rawaddr));
    if (bytes_sent) {
        logroot.infoStream() << "Sent " << msgtype(dhcpdata.type()) << " on iface " << req.ifindex
                             << ": xid: " << std::hex << dhcpdata.xid()
                             << " yiaddr: " << Tins::IPv4Address(iaddr.ipaddr).to_string()
                             <<" chaddr: " << req.chaddr.to_string();
    } else {
        logroot.error("Failed to send DHCP response on iface %i", req.ifindex);
    }
}

void process_dhcp() {
    // Init listening socket
    int res;
    int sockrcv = socket(AF_INET, SOCK_DGRAM, 0);
    struct sockaddr_in localaddr;
    memset((void *)&localaddr,0,sizeof(localaddr));
    localaddr.sin_family = AF_INET;
    localaddr.sin_port = htons(67);
    localaddr.sin_addr.s_addr = htonl(INADDR_ANY);

    int optval = 1;
    if (setsockopt(sockrcv, SOL_SOCKET, SO_BROADCAST, &optval, sizeof(optval)) < 0) {
        perror("Failed to set SO_BROADCAST on listening socket: ");
        exit(EXIT_FAILURE);
    }
    optval = 1;
    if (setsockopt(sockrcv, SOL_IP, IP_PKTINFO, &optval, sizeof(optval)) < 0) {
        perror("Failed to set SO_BROADCAST on listening socket: ");
        exit(EXIT_FAILURE);
    }

    if (bind(sockrcv, (struct sockaddr *) &localaddr, sizeof(localaddr)) < 0) {
        perror("Failed to bind recv socket on 0.0.0.0:67");
        exit(EXIT_FAILURE);
    }

    int sockraw = socket(AF_PACKET, SOCK_DGRAM, 0);

    // Receive loop
    while(true) {
        request req;
        res = recv_dhcp_packet(sockrcv, req);
        if (res < 0) {
            continue;
        }
        dhcp_send_reply(sockraw, req);
        //        break;
    }
    close(sockrcv);
}



int main(int argc, char *argv[])
{
    // Setup logging
    std::cout.setf(std::ios::unitbuf);
    log4cpp::Appender *appender_console = new log4cpp::OstreamAppender("console", &std::cout);
    log4cpp::PatternLayout *layout_console = new log4cpp::PatternLayout();
    layout_console->setConversionPattern("[%p] %m%n");
    appender_console->setLayout(layout_console);
    logroot.addAppender(appender_console);
    logroot.setPriority(log4cpp::Priority::INFO);

    logroot.infoStream() << "nasdhcp started";

    addr_table.active_index = 0;
    if (parse_config() < 0) {
        logroot.critStream() << "Failed to parse config!";
        return 1;
    }
    std::thread t_am(maintain_addr_map);
    process_dhcp();
    return 0;
}