#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <time.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <sys/time.h>
#include <errno.h>

#define NTP_TIMESTAMP_DELTA 2208988800ull
#define NTP_PORT 123
#define NTP_PACKET_SIZE 48

int main(int argc, char *argv[]) {
    if (argc != 2) {
        fprintf(stderr, "Usage: %s <ntp_server>\n", argv[0]);
        return 1;
    }

    const char *host = argv[1];
    int sockfd;
    struct sockaddr_in server_addr;
    struct hostent *server;

    // Create socket
    sockfd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sockfd < 0) {
        perror("socket");
        return 1;
    }

    // Set 5-second receive timeout
    struct timeval timeout;
    timeout.tv_sec = 5;
    timeout.tv_usec = 0;
    if (setsockopt(sockfd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout)) < 0) {
        perror("setsockopt");
        close(sockfd);
        return 1;
    }

    server = gethostbyname(host);
    if (!server) {
        fprintf(stderr, "Error: No such host %s\n", host);
        close(sockfd);
        return 1;
    }

    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(NTP_PORT);
    memcpy(&server_addr.sin_addr.s_addr, server->h_addr, server->h_length);

    // Create and initialize NTP packet
    unsigned char packet[NTP_PACKET_SIZE] = {0};
    packet[0] = 0b11100011;  // LI = 3 (unspecified), VN = 4, Mode = 3 (client)

    // Send the packet
    if (sendto(sockfd, packet, NTP_PACKET_SIZE, 0, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("sendto");
        close(sockfd);
        return 1;
    }

    // Receive the packet
    socklen_t addr_len = sizeof(server_addr);
    if (recvfrom(sockfd, packet, NTP_PACKET_SIZE, 0, (struct sockaddr *)&server_addr, &addr_len) < 0) {
        if (errno == EWOULDBLOCK || errno == EAGAIN) {
            fprintf(stderr, "Timeout: No response from server within 5 seconds.\n");
        } else {
            perror("recvfrom");
        }
        close(sockfd);
        return 1;
    }

    close(sockfd);

    // Extract the transmit timestamp (bytes 40-43)
    uint32_t transmit_time;
    memcpy(&transmit_time, &packet[40], sizeof(transmit_time));
    transmit_time = ntohl(transmit_time);

    time_t time_recv = transmit_time - NTP_TIMESTAMP_DELTA;

    // Set system time
    struct timeval tv = {
        .tv_sec = time_recv,
        .tv_usec = 0
    };

    if (settimeofday(&tv, NULL) < 0) {
        perror("settimeofday (are you root?)");
        return 1;
    }

    printf("System time set successfully to: %s", ctime(&tv.tv_sec));
    return 0;
}