#include <main.h>

// ⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿
// ⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⡟⠉⠀⢻⣿⣿⣿⣿⣿⣿⡿⠁⠀⠙⢿⣿⣿
// ⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⠏⠀⠀⠀⣾⣿⣿⣿⣿⣿⣿⠇⠀⠀⠀⠈⣿⣿
// ⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⠀⠀⠀⠀⠿⠛⠛⠛⠛⠿⣿⠀⠀⠀⠀⢀⣿⣿
// ⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⠀⣀⣤⡀⠀⠀⠀⠀⣀⡀⠀⠀⠀⢀⠀⢸⣿⣿
// ⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⡿⣴⣿⣿⠇⠀⠀⢠⣾⣿⣿⣆⠀⠀⠀⠀⣿⣿⣿
// ⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⡿⢡⣿⣿⣿⠀⠀⢀⣿⣿⣿⣿⣿⠀⠀⠀⢸⣿⣿⣿
// ⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⠁⢸⣿⡿⠃⠀⠀⢸⣿⣿⣿⣿⠏⠀⠀⠀⡘⣿⣿⣿
// ⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⡀⠀⠈⠀⠈⠭⠀⠈⠛⠿⠛⠋⠀⠀⠀⠀⠠⣿⣿⣿
// ⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣷⡑⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸⣿⣿
// ⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣷⠈⠂⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸⣿⣿
// ⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⡏⠀⢣⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸⣿⣿
// ⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⡿⠀⠀⠈⢃⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣼⣿⣿
// ⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣿⣿⣿
// ⣿⣿⣿⣿⣿⣿⣿⣿⡿⠃⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸⣿⣿⣿

const char* argp_program_version = "ft_ping 1.0";
const char* argp_program_bug_address = "<jrathelo@student.42nice.fr>";
static char doc[] = "Send ICMP ECHO_REQUEST packets to network hosts.\vMandatory or optional arguments to long options "
                    "are also mandatory or optional for any corresponding short options.\n\nOptions marked with (root "
                    "only) are available only to superuser.";
static char args_doc[] = "HOST ...";

#define ARG_TTL 256

static struct argp_option options[] = {
    {0, 0, 0, 0, "Options valid for all request types:\n", 0},
    {"verbose", 'v', 0, 0, "verbose output", 0},
    {"ttl", ARG_TTL, "N", 0, "specify N as time-to-live", 0},
    {0},
};

static error_t parse_opt(int key, char* arg, struct argp_state* state) {
    struct arguments* arguments = state->input;

    switch (key) {
        case 'v':
            arguments->verbose = 1;
            break;
        case ARG_TTL:
            arguments->ttl = strtoul(arg, NULL, 10);
            break;
        case ARGP_KEY_ARG:
            if (state->arg_num >= 2)
                argp_usage(state);
            arguments->args[state->arg_num] = arg;
            break;
        case ARGP_KEY_END:
            if (state->arg_num < 1)
                argp_usage(state);
            break;
        default:
            return ARGP_ERR_UNKNOWN;
    }
    return 0;
}

static struct argp argp = {options, parse_opt, args_doc, doc, NULL, NULL, NULL};

static bool stop = false;

int sock = 0;

static struct sockaddr_in destination = {0};

struct timespec sendt;

volatile unsigned long xmit = 0;
volatile unsigned long recvd = 0;
volatile unsigned long rept = 0;

uint16_t calc_checksum(void* b, size_t len) {
    unsigned short* buf = b;
    unsigned int sum = 0;
    unsigned short result;

    for (sum = 0; len > 1; len -= 2) {
        sum += *buf++;
    }

    if (len == 1) {
        sum += *(unsigned char*)buf;
    }

    sum = (sum >> 16) + (sum & 0xFFFF);
    sum += (sum >> 16);
    result = ~sum;
    return result;
}

void send_ping() {
    struct ping_pkt packet = {0};

    packet.hdr.type = ICMP_ECHO;
    packet.hdr.code = 0;
    packet.hdr.un.echo.id = getpid();
    packet.hdr.un.echo.sequence = xmit++;

    struct timeval tmp_tv;
    gettimeofday(&tmp_tv, NULL);

    memcpy(packet.msg, &tmp_tv, sizeof(struct timeval));
    for (long unsigned i = 0; i < sizeof(packet.msg) - sizeof(tmp_tv); ++i) {
        packet.msg[sizeof(tmp_tv) + i] = (char)i;
    }

    packet.hdr.checksum = calc_checksum(&packet, sizeof(packet));

    sendto(sock, &packet, sizeof(packet), 0, (struct sockaddr*)&destination, sizeof(destination));
    clock_gettime(CLOCK_MONOTONIC, &sendt);
    alarm(1);
}

void sig_handler(int sig) {
    if (sig == SIGINT) {
        stop = true;
    } else if (sig == SIGALRM) {
        send_ping();
    }
}

double nabs(double a) { return (a < 0) ? -a : a; }

double nsqrt(double a, double prec) {
    double b;
    double c;

    if (a < 0) {
        return 0;
    }
    if (a < prec) {
        return 0;
    }
    c = a / 2;
    do {
        b = c;
        c = (b + a / b) / 2;
    } while (nabs(c - b) > prec);

    return c;
}

void print_ip_hdr(int verbose, struct ip* ip_header) {
    if (verbose) {
        printf("IP Hdr dump:\n ");
        for (size_t i = 0; i < sizeof(*ip_header); i++) {
            printf("%02x", *((unsigned char*)ip_header + i));
            if (i % 2) {
                printf(" ");
            }
        }
        printf("\n");
    }

    printf("Vr HL TOS  Len   ID Flg  off TTL Pro  cks      Src\tDst\tData\n");
    printf(" %1x  %1x  %02x %04x %04x   %1x %04x  %02x  %02x %04x %s  %s ", ip_header->ip_v, ip_header->ip_hl,
           ip_header->ip_tos, (ip_header->ip_len > 0x2000) ? ntohs(ip_header->ip_len) : ip_header->ip_len,
           ntohs(ip_header->ip_id), (ntohs(ip_header->ip_off) & 0xe000) >> 13, ntohs(ip_header->ip_off) & 0x1fff,
           ip_header->ip_ttl, ip_header->ip_p, ntohs(ip_header->ip_sum),
           inet_ntoa(*((struct in_addr*)&ip_header->ip_src)), inet_ntoa(*((struct in_addr*)&ip_header->ip_dst)));
    unsigned char* cp = (unsigned char*)ip_header + sizeof(*ip_header);
    for (size_t l = ip_header->ip_hl << 2; l > sizeof(*ip_header); l--) {
        printf("%02x", *cp++);
    }
    printf("\n");
}

int main(int argc, char** argv) {
    struct arguments arguments;
    int timeout = 3000;
    struct timespec recvt;
    struct addrinfo hints = {0};
    struct addrinfo* res = NULL;
    struct pollfd fds[1];
    double min = INT32_MAX;
    double max = 0;
    double avg = 0;
    double stddev = 0;

    hints.ai_family = AF_INET;
    hints.ai_socktype = SOCK_STREAM;

    arguments.verbose = 0;
    arguments.ttl = 64;
    argp_parse(&argp, argc, argv, 0, 0, &arguments);

    // memset(&destination, 0, sizeof(destination));

    char ip_str[INET_ADDRSTRLEN] = {0};

    int status = getaddrinfo(arguments.args[0], NULL, &hints, &res);
    if (status != 0) {
        destination.sin_addr.s_addr = inet_addr(arguments.args[0]);
        if (destination.sin_addr.s_addr == INADDR_NONE) {
            dprintf(STDERR_FILENO, "%s: unknown host\n", argv[0]);
            exit(1);
        }
        destination.sin_family = AF_INET;
        destination.sin_port = htons(IPPORT_ECHO);
        memcpy(ip_str, arguments.args[0], INET_ADDRSTRLEN);
    } else {
        memcpy(&destination, res->ai_addr, sizeof(struct sockaddr_in));
        destination.sin_family = AF_INET;
        destination.sin_port = htons(IPPORT_ECHO);
        inet_ntop(res->ai_family, &destination.sin_addr, ip_str, INET_ADDRSTRLEN);
    }

    struct protoent* proto = getprotobyname("icmp");
    if (!proto) {
        dprintf(STDERR_FILENO, "%s: unknown protocol icmp.\n", argv[0]);
        exit(1);
    }

    sock = socket(AF_INET, SOCK_RAW, proto->p_proto);

    if (sock < 0) {
        if (errno == EPERM || errno == EACCES || errno == EPROTONOSUPPORT) {
            dprintf(STDERR_FILENO, "%s: Lacking privilege for icmp socket.\n", argv[0]);
        } else {
            dprintf(STDERR_FILENO, "%s: %s\n", argv[0], strerror(errno));
        }
        exit(1);
    }

    if (setsockopt(sock, IPPROTO_IP, IP_TTL, &arguments.ttl, sizeof(arguments.ttl)) < 0) {
        dprintf(STDERR_FILENO, "%s: setsockopt() failed: %s\n", argv[0], strerror(errno));
        close(sock);
        exit(1);
    }

    fds[0].fd = sock;
    fds[0].events = POLLIN;

    struct sigaction sig_handle = {0};

    sig_handle.sa_handler = sig_handler;

    sigaction(SIGINT, &sig_handle, 0);
    sigaction(SIGALRM, &sig_handle, 0);

    printf("PING %s (%s): %ld data bytes", arguments.args[0], ip_str, 64 - sizeof(struct icmphdr));
    if (arguments.verbose) {
        printf(", id 0x%04x = %u", getpid(), getpid());
    }
    printf("\n");

    send_ping();

    struct sockaddr_in localhost;
    inet_pton(AF_INET, "127.0.0.1", &(localhost.sin_addr));

    while (!stop) {
        int ret = poll(fds, 1, timeout);
        if (ret > 0) {
            if (fds[0].revents & POLLIN) {
                char buf[100] = {0};

                long int bytes_received = recvfrom(sock, buf, 100, 0, NULL, NULL);

                if (bytes_received < 0) {
                    perror("recvfrom failed");
                    return 1;
                }

                if ((size_t)bytes_received < sizeof(struct ip) + sizeof(struct icmphdr) + 8) {
                    dprintf(STDERR_FILENO, "packet too short (%ld bytes) from %s\n", bytes_received, arguments.args[0]);
                    if (!stop)
                        pause();
                    continue;
                }

                clock_gettime(CLOCK_MONOTONIC, &recvt);

                struct ip* ip_header = (struct ip*)buf;

                char strip[INET_ADDRSTRLEN] = {0};
                inet_ntop(AF_INET, &(ip_header->ip_src.s_addr), strip, INET_ADDRSTRLEN);
                struct icmphdr* icmp_reply = (struct icmphdr*)(buf + (ip_header->ip_hl << 2));
                if (icmp_reply->type == ICMP_ECHO && ip_header->ip_src.s_addr == localhost.sin_addr.s_addr) {
                    continue;
                } else if (icmp_reply->type != ICMP_ECHOREPLY) {
                    print_ip_hdr(arguments.verbose, ip_header);
                    continue;
                }

                // if (icmp_reply->code) {}

                int old_checksum = icmp_reply->checksum;
                icmp_reply->checksum = 0;

                if (old_checksum != calc_checksum(icmp_reply, bytes_received - (ip_header->ip_hl << 2))) {
                    dprintf(STDERR_FILENO, "checksum mismatch from %s\n", strip);
                    rept++;
                    if (!stop)
                        pause();
                    continue;
                }

                icmp_reply->checksum = old_checksum;

                double tmp =
                    ((recvt.tv_sec * 1000 + recvt.tv_nsec / 1000000) - (sendt.tv_sec * 1000 + sendt.tv_nsec / 1000000));
                tmp += ((double)((recvt.tv_nsec - sendt.tv_nsec) % 1000000) / 1000000.0);

                printf("%zd bytes from %s: icmp_seq=%d ttl=%d time=%.3f ms %s\n", bytes_received - sizeof(struct ip),
                       strip, icmp_reply->un.echo.sequence, ip_header->ip_ttl, tmp,
                       icmp_reply->un.echo.sequence < xmit ? "" : " (DUP!)");

                min = (((min) < (tmp)) ? (min) : (tmp));
                max = (((max) > (tmp)) ? (max) : (tmp));
                avg += tmp;
                stddev += tmp * tmp;
                if (icmp_reply->un.echo.sequence < recvd) {
                    rept++;
                } else {
                    recvd++;
                }
            }
        } else if (ret == 0) {
            printf("No reply received within the timeout period\n");
        } else {
            if (errno != EINTR) {
                perror("Poll failed");
            }
        }

        if (!stop)
            pause();
    }
    fflush(stdout);
    printf("---- %s ping statistics ----\n", arguments.args[0]);
    printf("%zu packets transmitted, %zu packets received, ", xmit, recvd);
    if (rept) {
        printf("+%zu duplicates, ", rept);
    }
    if (xmit) {
        if (recvd > xmit) {
            printf("-- somebody is printing forged packets!\n");
        } else {
            printf("%zu%% packet loss\n", (xmit - recvd) / xmit * 100);
        }
    }
    double total = recvd + rept;
    if (total > 0) {
        avg /= total;
        stddev = stddev / total - avg * avg;
        printf("round-trip min/avg/max/stddev = %.3f/%.3f/%.3f/%.3f ms\n", min, max, avg, nsqrt(stddev, 0.0005));
    }
    close(sock);
}