#include <algorithm>
#include <arpa/inet.h>
#include <boost/algorithm/hex.hpp>
#include <boost/algorithm/string.hpp>
#include <errno.h>
#include <iomanip>
#include <iostream>
#include <map>
#include <netinet/in.h>
#include <sstream>
#include <stdexcept>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <vector>

using namespace std;

const string name_server = "114.114.114.114";

struct unix_error : runtime_error {
    explicit unix_error(const string &s) : runtime_error(s + ": " + strerror(errno)) {}
};

struct host_error : exception {
    int rcode;
    string reason;
    explicit host_error(int rcode);
    const char *what() const noexcept override;
};

map<int, string> host_error_map = {
    {0, "NOERROR"},
    {1, "FORMERR"},
    {2, "SERVFAIL"},
    {3, "NXDOMAIN"},
};

host_error::host_error(int rcode) : rcode(rcode)
{
    string s;
    try {
        s = host_error_map.at(rcode);
    } catch (...) {
        s = "UNKNOWN";
    }
    reason = to_string(rcode) + "(" + s + ")";
}

const char *host_error::what() const noexcept
{
    return reason.c_str();
}

enum rr_type {
    A = 1,
    NS = 2,
    CNAME = 5,
    PTR = 12,
    HINFO = 13,
    MX = 15,
};

enum rr_class {
    IN = 1,
};

struct {
    rr_type type;
    string name;
} options;

struct __attribute__((__packed__)) dnshdr {
    unsigned short id;
    struct {
        unsigned short rd:1, tc:1, aa:1, opcode:4, qr:1, rcode:4, zero:3, ra:1;
    } flags;
    unsigned short qdcount, ancount, nscount, arcount;
};

struct question {
    string qname;
    unsigned short qtype, qclass;
    explicit question(const string &name, unsigned short type);
    explicit question(const unsigned char *hdr, const unsigned char *ptr);
    string data() const;
};

struct resource_record {
    string dname;
    unsigned short type, cls;
    unsigned int ttl;
    unsigned short rdatalen;
    string rdata;
    explicit resource_record(const unsigned char *hdr, const unsigned char *ptr);
};

struct dnsmsg {
    dnshdr dns_hdr;
    vector<question> questions;
    vector<resource_record> answers;
    vector<resource_record> authority;
    vector<resource_record> additional;
    dnsmsg(const string &name, unsigned short type);
    explicit dnsmsg(const unsigned char *hdr);
    string data() const;
};

string to_name(const unsigned char *hdr, const unsigned char *ptr)
{
    string name;
    while (*ptr) {
        if (*ptr >= 0xc0) {
            unsigned short off;
            memcpy(&off, ptr, sizeof(off));
            off = ntohs(off);
            off &= 0x3fff;
            ptr = hdr + off;
        } else {
            size_t n = *ptr;
            ++ptr;
            name += string((char *)ptr, n);
            name += ".";
            ptr += n;
        }
    }
    return name;
}

const unsigned char *ptr_after_name(const unsigned char *ptr)
{
    for (;;) {
        if (*ptr >= 0xc0) {
            ptr += 2;
            break;
        }
        if (*ptr == 0) {
            ++ptr;
            break;
        }
        size_t n = *ptr;
        ptr += (n + 1);
    }
    return ptr;
}

const unsigned char *ptr_after_question(const unsigned char *ptr)
{
    ptr = ptr_after_name(ptr);
    ptr += 4;
    return ptr;
}

const unsigned char *ptr_after_resource_record(const unsigned char *ptr)
{
    ptr = ptr_after_name(ptr);
    ptr += 8;
    unsigned short rdatalen;
    memcpy(&rdatalen, ptr, sizeof(rdatalen));
    rdatalen = ntohs(rdatalen);
    ptr += sizeof(rdatalen);
    ptr += rdatalen;
    return ptr;
}

bool is_ipv4(const string &s)
{
    return inet_addr(s.c_str()) != INADDR_NONE;
}

question::question(const string &name, unsigned short type)
    : qtype(htons(type)), qclass(htons(rr_class::IN))
{
    if (is_ipv4(name)) {
        vector<string> labels;
        boost::split(labels, name, boost::is_any_of("."));
        qname = labels[3] + "." +
                labels[2] + "." +
                labels[1] + "." +
                labels[0] + "." +
                "in-addr.arpa.";
        if (qtype == 0)
            qtype = htons(rr_type::PTR);
    } else {
        qname = name;
        if (!qname.empty() && qname[qname.size() - 1] != '.')
            qname += ".";
        if (qtype == 0)
            qtype = htons(rr_type::A);
    }
}

question::question(const unsigned char *hdr, const unsigned char *ptr)
{
    qname = to_name(hdr, ptr);
    ptr = ptr_after_name(ptr);
    memcpy(&qtype, ptr, sizeof(qtype));
    ptr += sizeof(qtype);
    memcpy(&qclass, ptr, sizeof(qclass));
}

string question::data() const
{
    string s;
    vector<string> labels;
    boost::split(labels, qname, boost::is_any_of("."));
    for (const string &label : labels) {
        s.push_back(label.size());
        s.append(label);
        if (label.empty())
            break;
    }
    s.append((char *)&qtype, sizeof(qtype));
    s.append((char *)&qclass, sizeof(qtype));
    return s;
}

resource_record::resource_record(const unsigned char *hdr, const unsigned char *ptr)
{
    dname = to_name(hdr, ptr);
    ptr = ptr_after_name(ptr);
    memcpy(&type, ptr, sizeof(type));
    ptr += sizeof(type);
    memcpy(&cls, ptr, sizeof(cls));
    ptr += sizeof(cls);
    memcpy(&ttl, ptr, sizeof(ttl));
    ptr += sizeof(ttl);
    memcpy(&rdatalen, ptr, sizeof(rdatalen));
    ptr += sizeof(rdatalen);
    switch (ntohs(type)) {
    case rr_type::A: {
        struct in_addr ip;
        memcpy(&ip, ptr, sizeof(ip));
        rdata = inet_ntoa(ip);
        break;
    }
    case rr_type::NS:
    case rr_type::CNAME:
    case rr_type::PTR:
    case rr_type::HINFO:
        rdata = to_name(hdr, ptr);
        break;
    case rr_type::MX: {
        unsigned short pref;
        memcpy(&pref, ptr, sizeof(pref));
        pref = ntohs(pref);
        ptr += sizeof(pref);
        rdata = to_string(pref) + " " + to_name(hdr, ptr);
        break;
    }
    default:
        break;
    }
}

dnsmsg::dnsmsg(const string &name, unsigned short type)
    : dns_hdr()
{
    dns_hdr.id = (unsigned short)getpid();
    dns_hdr.flags.rd = 1;       // recursion desired
    dns_hdr.qdcount = htons(1); // number of questions
    questions.push_back(question(name, type));
}

dnsmsg::dnsmsg(const unsigned char *hdr)
{
    auto *ptr = hdr;
    memcpy(&dns_hdr, ptr, sizeof(dns_hdr));
    if (dns_hdr.flags.rcode)
        throw host_error(dns_hdr.flags.rcode);
    ptr += sizeof(dns_hdr);
    for (int i = 0; i < ntohs(dns_hdr.qdcount); ++i) {
        question q(hdr, ptr);
        questions.push_back(q);
        ptr = ptr_after_question(ptr);
    }
    for (int i = 0; i < ntohs(dns_hdr.ancount); ++i) {
        resource_record rr(hdr, ptr);
        answers.push_back(rr);
        ptr = ptr_after_resource_record(ptr);
    }
}

string dnsmsg::data() const
{
    string s;
    s.append((char *)&dns_hdr, sizeof(dns_hdr));
    for (const question &q : questions)
        s.append(q.data());
    return s;
}

string class_name(unsigned short cls)
{
    switch (cls) {
    case rr_class::IN:
        return "IN";
    default:
        return "Unknown";
    }
}

map<rr_type, string> type_name_map = {
    {rr_type::A, "A"},
    {rr_type::NS, "NS"},
    {rr_type::CNAME, "CNAME"},
    {rr_type::PTR, "PTR"},
    {rr_type::HINFO, "HINFO"},
    {rr_type::MX, "MX"},
};

string type_name(unsigned short type)
{
    try {
        return type_name_map.at((rr_type)type);
    } catch (...) {
        return "Unknown";
    }
}

ostream &operator<<(ostream &out, const dnsmsg &rhs)
{
    if (rhs.answers.empty()) {
        out << "no record";
        return out;
    }

    auto it = max_element(rhs.answers.cbegin(), rhs.answers.cend(),
                          [](const resource_record &lhs, const resource_record &rhs) { return lhs.dname.size() < rhs.dname.size(); });
    for (int i = 0; i < ntohs(rhs.dns_hdr.ancount); ++i) {
        out << left
            << setw(it->dname.size() + 8) << rhs.answers[i].dname
            << setw(8) << ntohl(rhs.answers[i].ttl)
            << setw(8) << class_name(ntohs(rhs.answers[i].cls))
            << setw(8) << type_name(ntohs(rhs.answers[i].type))
            << setw(8) << rhs.answers[i].rdata
            << right;
        if (i != ntohs(rhs.dns_hdr.ancount) - 1)
            out << endl;
    }
    return out;
}

string hex(const void *buf, size_t len)
{
    string s;
    unsigned char *ptr = (unsigned char *)buf;
    boost::algorithm::hex_lower(ptr, ptr + len, back_inserter(s));
    return s;
}

void mydig()
{
    int sockfd = socket(AF_INET, SOCK_DGRAM, 0);
    if (sockfd < 0)
        throw unix_error("socket");

    struct sockaddr_in addr = {};
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = inet_addr(name_server.c_str());
    addr.sin_port = htons(53);

    dnsmsg qry(options.name, options.type);
    string d = qry.data();
    if (sendto(sockfd, d.data(), d.size(), 0, (struct sockaddr *)&addr, sizeof(addr)) < 0)
        throw unix_error("sendto");

    unsigned char buf[1024];
    ssize_t nrecv = recvfrom(sockfd, buf, sizeof(buf), 0, NULL, NULL);
    if (nrecv < 0)
        throw unix_error("recvfrom");
    dnsmsg rep(buf);
    cout << rep << endl;
}

map<string, rr_type> type_value_map = {
    {"A", rr_type::A},
    {"NS", rr_type::NS},
    {"CNAME", rr_type::CNAME},
    {"PTR", rr_type::PTR},
    {"HINFO", rr_type::HINFO},
    {"MX", rr_type::MX},
};

void parse_options(int argc, char *argv[])
{
    for (;;) {
        int c = getopt(argc, argv, ":t:");
        switch (c) {
        case -1:
            goto out;
        case 't': {
            string type = optarg;
            boost::to_upper(type);
            options.type = type_value_map.at(type);
            break;
        }
        default:
            throw invalid_argument("");
        }
    }
out:
    if (optind + 1 != argc)
        throw invalid_argument("");
    options.name = argv[optind++];
}

int main(int argc, char *argv[])
{
    try {
        parse_options(argc, argv);
        mydig();
    } catch (const logic_error &e) {
        cout << "usage: mydig [-t type] name" << endl;
    } catch (const exception &e) {
        cout << e.what() << endl;
    }
}