#include <stdio.h>
#include <string.h>
#include "pico/stdlib.h"
#include "pico/cyw43_arch.h"
#include "lwip/pbuf.h"
#include "lwip/udp.h"

#define WIFI_SSID "Wokwi-GUEST"
#define WIFI_PASSWORD ""
#define NTP_SERVER "time.google.com"
#define NTP_PORT 123

static struct udp_pcb *udp_client_pcb;

// NTP packet structure: 48 bytes
uint8_t ntp_request_packet[48] = {0x1B};

// Callback for when data is received
void udp_recv_callback(void *arg, struct udp_pcb *upcb, struct pbuf *p, const ip_addr_t *addr, u16_t port) {
    if (p == NULL) {
        printf("Connection closed.\n");
        return;
    }

    if (p->len >= 48) {
        // Extract the timestamp from the received NTP packet (bytes 40-43 for seconds since 1900)
        uint32_t seconds_since_1900;
        pbuf_copy_partial(p, &seconds_since_1900, sizeof(seconds_since_1900), 40);

        // Convert to host endianness
        seconds_since_1900 = ntohl(seconds_since_1900);

        // Convert to Unix epoch (subtract seconds from 1900 to 1970)
        const uint32_t NTP_UNIX_EPOCH_DIFF = 2208988800U;
        uint32_t unix_time = seconds_since_1900 - NTP_UNIX_EPOCH_DIFF;

        // Print the timestamp
        printf("NTP Time: %u (Unix Time)\n", unix_time);
    } else {
        printf("Received packet too small for NTP.\n");
    }

    pbuf_free(p);
}

// DNS callback function
static void dns_callback(const char *name, const ip_addr_t *ipaddr, void *callback_arg) {
    if (ipaddr == NULL) {
        printf("Failed to resolve DNS for %s\n", name);
        return;
    }

    printf("Resolved %s to %s\n", name, ipaddr_ntoa(ipaddr));

    // Send the NTP request packet
    struct pbuf *p = pbuf_alloc(PBUF_TRANSPORT, sizeof(ntp_request_packet), PBUF_RAM);
    if (!p) {
        printf("Failed to allocate pbuf.\n");
        return;
    }

    memcpy(p->payload, ntp_request_packet, sizeof(ntp_request_packet));
    udp_sendto(udp_client_pcb, p, ipaddr, NTP_PORT);
    pbuf_free(p);
}

void wifi_setup() {
    if (cyw43_arch_init()) {
        printf("Failed to initialize Wi-Fi module.\n");
        return;
    }

    cyw43_arch_enable_sta_mode();
    if (cyw43_arch_wifi_connect_timeout_ms(WIFI_SSID, WIFI_PASSWORD, CYW43_AUTH_WPA2_AES_PSK, 10000)) {
        printf("Cannot find SSID.\n");
        cyw43_arch_poll();
        cyw43_arch_deinit();
        sleep_ms(1000);
        wifi_setup();
    }
    printf("WiFi connected successfully.\n");
}

void print_network_info() {
    struct netif *netif = netif_list;
    printf("Network interface: %c%c\n", netif->name[0], netif->name[1]);
    printf("IP Address: %s\n", ip4addr_ntoa(netif_ip4_addr(netif)));
    printf("Netmask: %s\n", ip4addr_ntoa(netif_ip4_netmask(netif)));
    printf("Gateway: %s\n", ip4addr_ntoa(netif_ip4_gw(netif)));
}

void connect_to_ntp_server() {
    udp_client_pcb = udp_new();
    if (!udp_client_pcb) {
        printf("Failed to create UDP PCB.\n");
        return;
    }

    printf("Resolving %s...\n", NTP_SERVER);
    ip_addr_t server_ip;
    err_t err = dns_gethostbyname(NTP_SERVER, &server_ip, dns_callback, NULL);
    if (err == ERR_OK) {
        printf("DNS resolved immediately: %s\n", ipaddr_ntoa(&server_ip));

        // Send the NTP request packet immediately
        struct pbuf *p = pbuf_alloc(PBUF_TRANSPORT, sizeof(ntp_request_packet), PBUF_RAM);
        if (!p) {
            printf("Failed to allocate pbuf.\n");
            return;
        }

        memcpy(p->payload, ntp_request_packet, sizeof(ntp_request_packet));
        udp_sendto(udp_client_pcb, p, &server_ip, NTP_PORT);
        pbuf_free(p);
    } else if (err != ERR_INPROGRESS) {
        printf("DNS resolution failed: %d\n", err);
    }

    udp_recv(udp_client_pcb, udp_recv_callback, NULL);
}

int main() {
    stdio_init_all();
    wifi_setup();
    print_network_info();

    while (1) {
        connect_to_ntp_server();
        sleep_ms(10000); // Wait before reconnecting
    }
}