#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/time.h>
#include <errno.h>
#include <signal.h>
#include <pthread.h>
#include "udp_receiver.h"

// Debug flag for UDP receiver - commented by default
// Uncomment to see UDP debug messages
// #define DEBUG_UDP 1

// Host IP addresses (network 1, 2, 3)
static const char* HOST_IPS[] = {
    "10.10.2.1",    // Host 1
    "10.10.3.1",    // Host 2
    "10.10.4.1"     // Host 3
};

// Global variable for thread control
static int udp_running = 1;

// Host data storage with mutex protection
static HostData host_data[3] = {0};
static pthread_mutex_t host_mutex = PTHREAD_MUTEX_INITIALIZER;

// Parse UDP data in format: tcp,udp,icmp,rx,tx (COMMA separated)
static int parse_host_data(const char* buffer, HostData* data) {
    char temp_buffer[256];
    strncpy(temp_buffer, buffer, sizeof(temp_buffer) - 1);
    temp_buffer[sizeof(temp_buffer) - 1] = '\0';
    
    // Parse format: tcp,udp,icmp,rx,tx (COMMA separated)
    char *token;
    int field = 0;
    
    token = strtok(temp_buffer, ",");
    while (token != NULL && field < 5) {
        switch (field) {
            case 0: // tcp
                data->tcp = atol(token);
                break;
            case 1: // udp
                data->udp = atol(token);
                break;
            case 2: // icmp
                data->icmp = atol(token);
                break;
            case 3: // rx
                data->rx = atof(token);
                break;
            case 4: // tx
                data->tx = atof(token);
                break;
        }
        field++;
        token = strtok(NULL, ",");
    }
    
    gettimeofday(&data->last_update, NULL);
    
    return (field == 5) ? 0 : -1;
}

// Find host ID by IP address (returns 1-3, or 0 if not found)
static int get_host_id_by_ip(const char* ip) {
    for (int i = 0; i < 3; i++) {
        if (strcmp(ip, HOST_IPS[i]) == 0) {
            return i + 1; // Return 1, 2, or 3
        }
    }
    return 0; // Not a known host
}

// Get host data (thread-safe) - EXPORTED FUNCTION
int get_host_data(int host_id, HostData *data) {
    if (host_id < 1 || host_id > 3) {
        return -1;
    }
    
    pthread_mutex_lock(&host_mutex);
    
    // Check if data is stale (older than 10 seconds)
    struct timeval now;
    gettimeofday(&now, NULL);
    long age = (now.tv_sec - host_data[host_id-1].last_update.tv_sec);
    
    if (age > 10) {
        // Data is stale, return zeros
        memset(data, 0, sizeof(HostData));
        data->last_update = now;
    } else {
        // Copy fresh data
        memcpy(data, &host_data[host_id-1], sizeof(HostData));
    }
    
    pthread_mutex_unlock(&host_mutex);
    return 0;
}

void* udp_receiver_thread(void* arg) {
    (void)arg;  // Explicitly mark as unused
    int sockfd;
    struct sockaddr_in server_addr, client_addr;
    socklen_t client_len;
    char buffer[1024];
    
    // Create UDP socket
    sockfd = socket(AF_INET, SOCK_DGRAM, 0);
    if (sockfd < 0) {
        perror("UDP socket creation failed");
        return NULL;
    }
    
    // Set up server address
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(UDP_RECEIVER_PORT);  // Using 9998 as original
    
    // Bind socket
    if (bind(sockfd, (const struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("UDP bind failed");
        close(sockfd);
        return NULL;
    }
    
    while (udp_running) {
        client_len = sizeof(client_addr);
        memset(buffer, 0, sizeof(buffer));
        
        // Receive data with timeout to check running flag
        struct timeval tv;
        fd_set readfds;
        
        FD_ZERO(&readfds);
        FD_SET(sockfd, &readfds);
        tv.tv_sec = 1;
        tv.tv_usec = 0;
        
        int ret = select(sockfd + 1, &readfds, NULL, NULL, &tv);
        
        if (ret > 0 && FD_ISSET(sockfd, &readfds)) {
            int n = recvfrom(sockfd, buffer, sizeof(buffer) - 1, 0,
                            (struct sockaddr *)&client_addr, &client_len);
            
            if (n > 0) {
                // Remove newline if present
                buffer[strcspn(buffer, "\n")] = 0;
                
                // Get client IP
                char client_ip[INET_ADDRSTRLEN];
                inet_ntop(AF_INET, &client_addr.sin_addr, client_ip, INET_ADDRSTRLEN);
                
                // Check if this is from a known host
                int host_id = get_host_id_by_ip(client_ip);
                
                if (host_id > 0) {
                    // Parse host data
                    HostData new_data;
                    if (parse_host_data(buffer, &new_data) == 0) {
                        // Store in appropriate host slot (thread-safe)
                        pthread_mutex_lock(&host_mutex);
                        memcpy(&host_data[host_id-1], &new_data, sizeof(HostData));
                        pthread_mutex_unlock(&host_mutex);
                        
                        // Display received data only if DEBUG_UDP is defined
#ifdef DEBUG_UDP
                        printf("UDP_RECV: Host %d (%s) - TCP:%ld UDP:%ld ICMP:%ld RX:%.2f TX:%.2f\n",
                               host_id, client_ip, 
                               new_data.tcp, new_data.udp, new_data.icmp,
                               new_data.rx, new_data.tx);
                        fflush(stdout);
#endif
                    } else {
#ifdef DEBUG_UDP
                        printf("UDP_RECV: ERROR parsing data from %s: %s\n", client_ip, buffer);
                        fflush(stdout);
#endif
                    }
                } else {
#ifdef DEBUG_UDP
                    printf("UDP_RECV: Unknown host %s: %s\n", client_ip, buffer);
                    fflush(stdout);
#endif
                }
            }
        }
    }
    
    close(sockfd);
    return NULL;
}

void start_udp_receiver(pthread_t *thread_id) {
    udp_running = 1;
    if (pthread_create(thread_id, NULL, udp_receiver_thread, NULL) != 0) {
        fprintf(stderr, "WARNING: Failed to create UDP receiver thread\n");
    }
}

void stop_udp_receiver(pthread_t thread_id) {
    udp_running = 0;
    pthread_join(thread_id, NULL);
    pthread_mutex_destroy(&host_mutex);
}