#include "test_udp_socket.h"
#include "../lib/debug_config.h"
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <errno.h>
#include <fcntl.h>
#include <pthread.h>

#ifndef DEBUG_CATEGORY_TEST
#define DEBUG_CATEGORY_TEST (1 << 30)
#endif

// Global socket registry for fd-based operations
static struct test_udp_socket* g_socket_registry[1024] = {NULL};
pthread_mutex_t g_registry_mutex = PTHREAD_MUTEX_INITIALIZER;  // Make it non-static for external access
static int g_next_fd = 1000; // Start with high fd numbers to avoid conflicts

// Helper function to allocate packet
static struct test_udp_packet* allocate_packet(const uint8_t* data, size_t len,
                                               const struct sockaddr* addr, socklen_t addr_len) {
    struct test_udp_packet* packet = calloc(1, sizeof(struct test_udp_packet));
    if (!packet) return NULL;
    
    packet->data = malloc(len);
    if (!packet->data) {
        free(packet);
        return NULL;
    }
    
    memcpy(packet->data, data, len);
    packet->len = len;
    
    if (addr && addr_len > 0 && addr_len <= sizeof(struct sockaddr_storage)) {
        memcpy(&packet->addr, addr, addr_len);
        packet->addr_len = addr_len;
    } else {
        packet->addr_len = 0;
    }
    
    return packet;
}

// Helper function to free packet
static void free_packet(struct test_udp_packet* packet) {
    if (!packet) return;
    if (packet->data) free(packet->data);
    free(packet);
}

// Create virtual UDP socket
struct test_udp_socket* test_udp_socket_create(int family) {
    if (family != AF_INET && family != AF_INET6) {
        DEBUG_ERROR(DEBUG_CATEGORY_TEST, "Unsupported address family: %d", family);
        return NULL;
    }
    
    struct test_udp_socket* sock = calloc(1, sizeof(struct test_udp_socket));
    if (!sock) {
        DEBUG_ERROR(DEBUG_CATEGORY_TEST, "Failed to allocate UDP socket");
        return NULL;
    }
    
    // Assign virtual file descriptor
    pthread_mutex_lock(&g_registry_mutex);
    sock->fd = g_next_fd++;
    if (g_next_fd >= 2000) g_next_fd = 1000; // Wrap around
    pthread_mutex_unlock(&g_registry_mutex);
    
    sock->family = family;
    sock->bound = false;
    sock->nonblocking = false;
    sock->recv_queue_head = NULL;
    sock->recv_queue_tail = NULL;
    sock->recv_queue_size = 0;
    
    // Register socket
    test_udp_socket_register(sock);
    
    DEBUG_INFO(DEBUG_CATEGORY_TEST, "Created virtual UDP socket: fd=%d, family=%d", sock->fd, family);
    return sock;
}

// Destroy virtual UDP socket
void test_udp_socket_destroy(struct test_udp_socket* sock) {
    if (!sock) return;
    
    DEBUG_INFO(DEBUG_CATEGORY_TEST, "Destroying virtual UDP socket: fd=%d (stats: sent=%zu/%zu, recv=%zu/%zu, errors=%zu/%zu)",
               sock->fd, sock->stats.packets_sent, sock->stats.bytes_sent,
               sock->stats.packets_received, sock->stats.bytes_received,
               sock->stats.send_errors, sock->stats.recv_errors);
    
    // Unregister socket
    test_udp_socket_unregister(sock);
    
    // Free receive queue
    struct test_udp_packet* packet = sock->recv_queue_head;
    while (packet) {
        struct test_udp_packet* next = packet->next;
        free_packet(packet);
        packet = next;
    }
    
    free(sock);
}

// Bind virtual socket
int test_udp_socket_bind(struct test_udp_socket* sock, const struct sockaddr* addr, socklen_t len) {
    if (!sock || !addr || len == 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_TEST, "Invalid parameters for socket bind");
        return -1;
    }
    
    if (sock->bound) {
        DEBUG_WARN(DEBUG_CATEGORY_TEST, "Socket already bound");
        return 0;
    }
    
    if (len > sizeof(struct sockaddr_storage)) {
        DEBUG_ERROR(DEBUG_CATEGORY_TEST, "Address too large: %d", len);
        return -1;
    }
    
    memcpy(&sock->local_addr, addr, len);
    sock->bound = true;
    
    DEBUG_INFO(DEBUG_CATEGORY_TEST, "Bound virtual UDP socket fd=%d", sock->fd);
    return 0;
}

// Set socket options (limited implementation for testing)
int test_udp_socket_setsockopt(struct test_udp_socket* sock, int level, int optname, 
                              const void *optval, socklen_t optlen) {
    if (!sock) return -1;
    
    // For testing purposes, just log the option setting
    DEBUG_DEBUG(DEBUG_CATEGORY_TEST, "setsockopt on fd=%d: level=%d, optname=%d, optlen=%d", 
                sock->fd, level, optname, optlen);
    
    // Handle some common options
    if (level == SOL_SOCKET) {
        switch (optname) {
            case SO_REUSEADDR:
            case SO_REUSEPORT:
                // Always allow for testing
                return 0;
            case SO_BROADCAST:
                return 0;
            default:
                break;
        }
    }
    
    return 0; // Pretend success for most options
}

// Get socket options
int test_udp_socket_getsockopt(struct test_udp_socket* sock, int level, int optname,
                              void *optval, socklen_t *optlen) {
    if (!sock || !optval || !optlen) return -1;
    
    // For testing purposes, return reasonable defaults
    if (level == SOL_SOCKET) {
        switch (optname) {
            case SO_ERROR:
                *(int*)optval = 0;
                *optlen = sizeof(int);
                return 0;
            default:
                break;
        }
    }
    
    return -1; // Not implemented
}

// Set non-blocking mode
int test_udp_socket_set_nonblocking(struct test_udp_socket* sock, bool nonblocking) {
    if (!sock) return -1;
    
    sock->nonblocking = nonblocking;
    DEBUG_DEBUG(DEBUG_CATEGORY_TEST, "Set nonblocking on fd=%d: %s", sock->fd, nonblocking ? "true" : "false");
    return 0;
}

// Send packet using virtual socket
ssize_t test_udp_socket_sendto(struct test_udp_socket* sock, const void *buf, size_t len, int flags,
                              const struct sockaddr *dest_addr, socklen_t addr_len) {
    if (!sock || !buf || len == 0) {
        return -1;
    }
    
    const uint8_t* data = (const uint8_t*)buf;
    
    // Call packet sent callback if available
    if (sock->packet_sent) {
        sock->packet_sent(data, len, dest_addr, addr_len, sock->context);
    }
    
    sock->stats.packets_sent++;
    sock->stats.bytes_sent += len;
    
    DEBUG_DEBUG(DEBUG_CATEGORY_TEST, "Sent packet on fd=%d: %zu bytes to family=%d", 
                sock->fd, len, dest_addr ? dest_addr->sa_family : -1);
    return len;
}

// Receive packet from virtual socket
ssize_t test_udp_socket_recvfrom(struct test_udp_socket* sock, void *buf, size_t len, int flags,
                                struct sockaddr *src_addr, socklen_t *addr_len) {
    if (!sock || !buf || len == 0) {
        return -1;
    }
    
    uint8_t* buffer = (uint8_t*)buf;
    
    // Check if we have packets in queue
    pthread_mutex_lock(&g_registry_mutex);
    struct test_udp_packet* packet = sock->recv_queue_head;
    if (!packet) {
        pthread_mutex_unlock(&g_registry_mutex);
        errno = EAGAIN;
        return -1;
    }
    
    // Remove packet from queue
    sock->recv_queue_head = packet->next;
    if (!sock->recv_queue_head) {
        sock->recv_queue_tail = NULL;
    }
    sock->recv_queue_size--;
    pthread_mutex_unlock(&g_registry_mutex);
    
    // Copy packet data
    size_t copy_len = packet->len < len ? packet->len : len;
    memcpy(buffer, packet->data, copy_len);
    
    // Copy source address if requested
    if (src_addr && addr_len && packet->addr_len > 0) {
        socklen_t copy_addr_len = packet->addr_len < *addr_len ? packet->addr_len : *addr_len;
        memcpy(src_addr, &packet->addr, copy_addr_len);
        *addr_len = copy_addr_len;
    }
    
    sock->stats.packets_received++;
    sock->stats.bytes_received += copy_len;
    
    DEBUG_DEBUG(DEBUG_CATEGORY_TEST, "Received packet on fd=%d: %zu bytes", sock->fd, copy_len);
    
    free_packet(packet);
    return copy_len;
}

// Inject packet into receive queue (simulates incoming packet)
int test_udp_socket_inject(struct test_udp_socket* sock, 
                          const uint8_t* data, size_t len,
                          const struct sockaddr* src_addr, socklen_t addr_len) {
    if (!sock || !data || len == 0) {
        return -1;
    }
    
    if (sock->recv_queue_size >= TEST_UDP_MAX_QUEUE_SIZE) {
        DEBUG_WARN(DEBUG_CATEGORY_TEST, "Receive queue full on fd=%d", sock->fd);
        return -1;
    }
    
    struct test_udp_packet* packet = allocate_packet(data, len, src_addr, addr_len);
    if (!packet) {
        return -1;
    }
    
    pthread_mutex_lock(&g_registry_mutex);
    if (sock->recv_queue_tail) {
        sock->recv_queue_tail->next = packet;
    } else {
        sock->recv_queue_head = packet;
    }
    sock->recv_queue_tail = packet;
    sock->recv_queue_size++;
    pthread_mutex_unlock(&g_registry_mutex);
    
    DEBUG_DEBUG(DEBUG_CATEGORY_TEST, "Injected packet into fd=%d: %zu bytes", sock->fd, len);
    return 0;
}

// Get virtual socket file descriptor
int test_udp_socket_get_fd(struct test_udp_socket* sock) {
    return sock ? sock->fd : -1;
}

// Get socket statistics
void test_udp_socket_get_stats(struct test_udp_socket* sock, size_t* packets_sent, 
                              size_t* packets_received, size_t* bytes_sent, 
                              size_t* bytes_received, size_t* send_errors, size_t* recv_errors) {
    if (!sock) return;
    
    if (packets_sent) *packets_sent = sock->stats.packets_sent;
    if (packets_received) *packets_received = sock->stats.packets_received;
    if (bytes_sent) *bytes_sent = sock->stats.bytes_sent;
    if (bytes_received) *bytes_received = sock->stats.bytes_received;
    if (send_errors) *send_errors = sock->stats.send_errors;
    if (recv_errors) *recv_errors = sock->stats.recv_errors;
}

// Reset socket statistics
void test_udp_socket_reset_stats(struct test_udp_socket* sock) {
    if (!sock) return;
    
    memset(&sock->stats, 0, sizeof(sock->stats));
    DEBUG_INFO(DEBUG_CATEGORY_TEST, "Reset statistics for virtual UDP socket fd=%d", sock->fd);
}

// Global virtual socket registry for fd-based operations
struct test_udp_socket* test_udp_socket_find_by_fd(int fd) {
    pthread_mutex_lock(&g_registry_mutex);
    struct test_udp_socket* sock = NULL;
    if (fd >= 0 && fd < 1024) {
        sock = g_socket_registry[fd];
    }
    pthread_mutex_unlock(&g_registry_mutex);
    return sock;
}

void test_udp_socket_register(struct test_udp_socket* sock) {
    if (!sock) return;
    
    pthread_mutex_lock(&g_registry_mutex);
    if (sock->fd >= 0 && sock->fd < 1024) {
        g_socket_registry[sock->fd] = sock;
    }
    pthread_mutex_unlock(&g_registry_mutex);
}

void test_udp_socket_unregister(struct test_udp_socket* sock) {
    if (!sock) return;
    
    pthread_mutex_lock(&g_registry_mutex);
    if (sock->fd >= 0 && sock->fd < 1024) {
        g_socket_registry[sock->fd] = NULL;
    }
    pthread_mutex_unlock(&g_registry_mutex);
}