#include <iostream>
#include <thread>
#include <dnet.h>
#include <cstring>
#include <bits/socket.h>
#include <unistd.h>

struct icmp_packet;

uint16_t checksum(const icmp_packet *packet);

struct icmp_packet {
    icmp_hdr header{};
    uint16_t id = getpid();
    uint16_t seq = 0;
    char msg[56]{};

    explicit icmp_packet(const char *msg) {
        header.icmp_type = ICMP_ECHO;
        header.icmp_code = 0;
        strncpy(this->msg, msg, 55);
        update();
    }

    inline void update() {
        header.icmp_cksum = 0; // Important: we need to clear previous checksum before calculating a new one
        header.icmp_cksum = checksum(this);
    }
};

uint16_t checksum(const icmp_packet *packet) {
    int len = sizeof(*packet);
    auto *current = reinterpret_cast<const uint16_t *>(packet);
    uint32_t sum = 0;
    while (len > 1) {
        len -= 2;
        sum += *current++;
    }
    if (len == 1)
        sum += *reinterpret_cast<const uint8_t *>(current);
    sum = (sum >> 16u) + (sum & 0xFFFFu);
    sum += sum >> 16u;
    return ~sum;
}

char *dns_lookup(const char *hostname, sockaddr_in &addr, bool &is_ipv6) {
    static char ip[INET6_ADDRSTRLEN];
    static addrinfo hint{}, *res = nullptr;
    hint.ai_family = PF_UNSPEC;
    hint.ai_flags = AI_ALL;

    if (getaddrinfo(hostname, nullptr, &hint, &res)) {
        std::cerr << "Invalid address" << std::endl;
        return nullptr;
    }
    is_ipv6 = res->ai_family == AF_INET6;
    memcpy(&addr, res->ai_addr, sizeof(addr));
    inet_ntop(res->ai_family, &addr.sin_addr, ip, res->ai_addrlen);
    return ip;
}

void send_request(int sock_fd, const sockaddr_in &addr, const char *ip_addr, bool is_ipv6 = false) {
    icmp_packet packet("echo request");
    icmp_packet buffer("echo reply");
    sockaddr r_addr{};
    int cnt = 4;
    while (cnt--) {
        if (sendto(sock_fd, &packet, sizeof(packet), 0,
                   reinterpret_cast<const struct sockaddr *>(&addr), sizeof(addr)) < 1) {
            std::cerr << "Failed to send packet" << std::endl;
            continue;
        }
        socklen_t r_len = 0;
        recvfrom(sock_fd, &buffer, sizeof(buffer), 0, &r_addr, &r_len);
        if (buffer.header.icmp_type != 69 || buffer.header.icmp_code) {
            std::cout << "Broken packet" << std::endl;
            continue;
        }
        ++packet.seq;
        packet.update();
        using namespace std::chrono_literals;
        std::this_thread::sleep_for(500ms);
    }

}

int main(int argc, char **argv) {
    if (argc < 2) {
        std::cout << "Usage: " << argv[0] << " " << "(hostname|ip address)" << std::endl;
        return 1;
    }
    bool is_ipv6 = false;
    sockaddr_in addr{};
    char *ip = dns_lookup(argv[1], addr, is_ipv6);
    puts(ip);
    int sock_fd = socket(is_ipv6 ? AF_INET6 : AF_INET, SOCK_RAW, is_ipv6 ? IPPROTO_ICMPV6 : IPPROTO_ICMP);
    if (sock_fd < 0) {
        std::cerr << "Cannot open socket fd\n" << "Consider running with sudo" << std::endl;
        return sock_fd;
    }
    int ttl = 64;
    setsockopt(sock_fd, is_ipv6 ? SOL_IPV6 : SOL_IP, IP_TTL, &ttl, sizeof(ttl));
    send_request(sock_fd, addr, ip, is_ipv6);
    return 0;
}