/**
 * @file control_socket.c
 * @brief Control socket implementation for statistics and diagnostics
 */

#include "control_socket.h"
#include "connection.h"
#include "etcp.h"
#include "pkt_normalizer.h"
#include "routing.h"
#include "tun_if.h"
#include "config_parser.h"
#include "utun_state.h"

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

// Internal control socket structure
struct control_socket {
    int sockfd;
    struct sockaddr_in addr;
    uint16_t port;
    void *app_state;              // Pointer to utun_state_t
    uint16_t sequence_counter;
    
    // Statistics counters (for cumulative stats)
    uint64_t total_requests;
    uint64_t total_errors;
    
    // Timestamps
    struct timeval start_time;
};

// Forward declarations
static int collect_statistics(control_socket_t *cs, control_stats_packet_t *stats);
static int collect_system_status(control_socket_t *cs, system_status_t *status);
static void get_etcp_stats(epkt_t *etcp, etcp_stats_t *etcp_stats);
static void get_pkt_normalizer_stats(pkt_normalizer_pair *pair, pkt_normalizer_stats_t *pn_stats);
static int handle_stats_request(control_socket_t *cs, struct sockaddr_in *client, uint16_t sequence);
static int handle_status_request(control_socket_t *cs, struct sockaddr_in *client, uint16_t sequence);
static int handle_reset_request(control_socket_t *cs, struct sockaddr_in *client, uint16_t sequence);
static int send_error_response(control_socket_t *cs, struct sockaddr_in *client, uint16_t sequence, control_err_t error);
static int send_response(control_socket_t *cs, struct sockaddr_in *client, void *data, size_t len);

// Create UDP control socket
control_socket_t* control_socket_create(const char *ip, uint16_t port) {
    control_socket_t *cs = calloc(1, sizeof(control_socket_t));
    if (!cs) {
        return NULL;
    }
    
    // Create UDP socket
    cs->sockfd = socket(AF_INET, SOCK_DGRAM, 0);
    if (cs->sockfd < 0) {
        perror("control_socket: socket");
        free(cs);
        return NULL;
    }
    
    // Set non-blocking
    int flags = fcntl(cs->sockfd, F_GETFL, 0);
    if (flags < 0) {
        perror("control_socket: fcntl GETFL");
        close(cs->sockfd);
        free(cs);
        return NULL;
    }
    if (fcntl(cs->sockfd, F_SETFL, flags | O_NONBLOCK) < 0) {
        perror("control_socket: fcntl SETFL");
        close(cs->sockfd);
        free(cs);
        return NULL;
    }
    
    // Enable address reuse
    int opt = 1;
    if (setsockopt(cs->sockfd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
        perror("control_socket: setsockopt SO_REUSEADDR");
        // Continue anyway
    }
    
    // Bind to address
    memset(&cs->addr, 0, sizeof(cs->addr));
    cs->addr.sin_family = AF_INET;
    cs->addr.sin_port = htons(port);
    
    if (ip && ip[0] != '\0') {
        if (inet_pton(AF_INET, ip, &cs->addr.sin_addr) != 1) {
            fprintf(stderr, "control_socket: invalid IP address: %s\n", ip);
            close(cs->sockfd);
            free(cs);
            return NULL;
        }
    } else {
        cs->addr.sin_addr.s_addr = INADDR_ANY;
    }
    
    if (bind(cs->sockfd, (struct sockaddr*)&cs->addr, sizeof(cs->addr)) < 0) {
        perror("control_socket: bind");
        close(cs->sockfd);
        free(cs);
        return NULL;
    }
    
    cs->port = port;
    cs->sequence_counter = 0;
    cs->total_requests = 0;
    cs->total_errors = 0;
    gettimeofday(&cs->start_time, NULL);
    
    printf("[CONTROL] Control socket listening on %s:%u\n", 
           ip ? ip : "0.0.0.0", port);
    
    return cs;
}

// Destroy control socket
void control_socket_destroy(control_socket_t *cs) {
    if (!cs) return;
    
    if (cs->sockfd >= 0) {
        close(cs->sockfd);
    }
    
    free(cs);
}

// Get socket file descriptor for polling
int control_socket_get_fd(const control_socket_t *cs) {
    return cs ? cs->sockfd : -1;
}

// Set application state pointer
void control_socket_set_state(control_socket_t *cs, void *state) {
    if (cs) {
        cs->app_state = state;
    }
}

// Process incoming requests
int control_socket_process(control_socket_t *cs, void *state) {
    if (!cs || !state) return -1;
    
    // Use provided state or stored state
    if (state) {
        cs->app_state = state;
    }
    
    struct sockaddr_in client_addr;
    socklen_t addr_len = sizeof(client_addr);
    control_request_packet_t request;
    
    // Receive request
    ssize_t n = recvfrom(cs->sockfd, &request, sizeof(request), 0,
                        (struct sockaddr*)&client_addr, &addr_len);
    
    if (n < 0) {
        if (errno == EAGAIN || errno == EWOULDBLOCK) {
            return 0; // No data available
        }
        perror("control_socket: recvfrom");
        return -1;
    }
    
    if (n != sizeof(request)) {
        fprintf(stderr, "control_socket: invalid request size: %zd\n", n);
        return 0;
    }
    
    cs->total_requests++;
    
    // Handle request based on command
    switch (request.command) {
        case CONTROL_CMD_GET_STATS:
            return handle_stats_request(cs, &client_addr, request.sequence);
            
        case CONTROL_CMD_GET_STATUS:
            return handle_status_request(cs, &client_addr, request.sequence);
            
        case CONTROL_CMD_RESET_STATS:
            return handle_reset_request(cs, &client_addr, request.sequence);
            
        default:
            return send_error_response(cs, &client_addr, request.sequence, 
                                      CONTROL_ERR_INVALID_CMD);
    }
}

// Handle statistics request
static int handle_stats_request(control_socket_t *cs, struct sockaddr_in *client, 
                               uint16_t sequence) {
    control_stats_packet_t stats_packet = {0};
    
    // Fill response header
    stats_packet.response_type = CONTROL_RESP_STATS;
    stats_packet.error_code = CONTROL_ERR_NONE;
    stats_packet.sequence = sequence;
    
    // Get current timestamp
    struct timeval tv;
    gettimeofday(&tv, NULL);
    stats_packet.timestamp = (uint64_t)tv.tv_sec * 1000000 + tv.tv_usec;
    
    // Collect statistics from all modules
    if (collect_statistics(cs, &stats_packet) < 0) {
        return send_error_response(cs, client, sequence, CONTROL_ERR_INTERNAL);
    }
    
    // Send response
    return send_response(cs, client, &stats_packet, sizeof(stats_packet));
}

// Handle status request
static int handle_status_request(control_socket_t *cs, struct sockaddr_in *client,
                                uint16_t sequence) {
    control_status_packet_t status_packet = {0};
    
    status_packet.response_type = CONTROL_RESP_STATUS;
    status_packet.error_code = CONTROL_ERR_NONE;
    status_packet.sequence = sequence;
    
    // Collect system status
    if (collect_system_status(cs, &status_packet.system) < 0) {
        status_packet.system.uptime_seconds = 0;
    }
    
    // Set running flag
    if (cs->app_state) {
        // Assume utun_state_t has 'running' field at same offset
        // We'll fix this when integrating with utun.c
        status_packet.running = 1;
    }
    
    snprintf(status_packet.status_msg, sizeof(status_packet.status_msg),
             "utun control socket active on port %u", cs->port);
    
    return send_response(cs, client, &status_packet, sizeof(status_packet));
}

// Handle reset statistics request
static int handle_reset_request(control_socket_t *cs, struct sockaddr_in *client,
                               uint16_t sequence) {
    // TODO: Implement statistics reset
    // For now, just send success response
    
    control_status_packet_t response = {0};
    response.response_type = CONTROL_RESP_STATUS;
    response.error_code = CONTROL_ERR_NONE;
    response.sequence = sequence;
    
    // Get basic system status
    collect_system_status(cs, &response.system);
    
    snprintf(response.status_msg, sizeof(response.status_msg),
             "Statistics reset requested (not implemented)");
    
    return send_response(cs, client, &response, sizeof(response));
}

// Send error response
static int send_error_response(control_socket_t *cs, struct sockaddr_in *client,
                              uint16_t sequence, control_err_t error) {
    control_status_packet_t response = {0};
    
    response.response_type = CONTROL_RESP_ERROR;
    response.error_code = error;
    response.sequence = sequence;
    
    const char *error_msg = "Unknown error";
    switch (error) {
        case CONTROL_ERR_INVALID_CMD:
            error_msg = "Invalid command";
            break;
        case CONTROL_ERR_INTERNAL:
            error_msg = "Internal error";
            break;
        default:
            break;
    }
    
    snprintf(response.status_msg, sizeof(response.status_msg), "%s", error_msg);
    
    cs->total_errors++;
    
    return send_response(cs, client, &response, sizeof(response));
}

// Send response to client
static int send_response(control_socket_t *cs, struct sockaddr_in *client,
                        void *data, size_t len) {
    ssize_t sent = sendto(cs->sockfd, data, len, 0,
                         (struct sockaddr*)client, sizeof(*client));
    
    if (sent != (ssize_t)len) {
        perror("control_socket: sendto");
        return -1;
    }
    
    return 0;
}

// Collect statistics from all modules
static int collect_statistics(control_socket_t *cs, control_stats_packet_t *stats) {
    if (!cs->app_state) {
        return -1;
    }
    
    utun_state_t *state = (utun_state_t*)cs->app_state;
    
    // Zero out stats
    memset(&stats->pkt_normalizer, 0, sizeof(stats->pkt_normalizer));
    memset(&stats->etcp, 0, sizeof(stats->etcp));
    memset(&stats->routing, 0, sizeof(stats->routing));
    memset(&stats->tun, 0, sizeof(stats->tun));
    stats->conn_count = 0;
    
    // Collect connection statistics
    if (state->connections && state->connection_count > 0) {
        int max_conn = state->connection_count;
        if (max_conn > MAX_CONN_STATS) max_conn = MAX_CONN_STATS;
        
        for (int i = 0; i < max_conn; i++) {
            conn_handle_t *conn = state->connections[i];
            if (!conn) continue;
            
            conn_stats_t conn_stats;
            if (conn_get_stats(conn, &conn_stats) == 0) {
                // Aggregate connection stats
                stats->etcp.bytes_sent_total += conn_stats.bytes_sent;
                stats->etcp.bytes_received_total += conn_stats.bytes_received;
                stats->etcp.retransmissions += conn_stats.retransmissions;
                stats->pkt_normalizer.fragments_assembled += conn_stats.fragments_assembled;
                // Note: etcp and pkt_normalizer stats are aggregated across connections
                // Additional per-connection stats could be stored in conn_stats_entry_t
            }
            
            // Fill connection entry (if we have array)
            if (i < MAX_CONN_STATS) {
                if (state->config && state->config->connections) {
                    strncpy(stats->connections[i].name, state->config->connections[i].name, 
                            sizeof(stats->connections[i].name) - 1);
                    stats->connections[i].name[sizeof(stats->connections[i].name) - 1] = '\0';
                } else {
                    snprintf(stats->connections[i].name, sizeof(stats->connections[i].name), "conn%d", i);
                }
                stats->connections[i].bytes_sent = conn_stats.bytes_sent;
                stats->connections[i].bytes_received = conn_stats.bytes_received;
                stats->connections[i].packets_sent = conn_stats.packets_sent;
                stats->connections[i].packets_received = conn_stats.packets_received;
                stats->connections[i].retransmissions = conn_stats.retransmissions;
                stats->connections[i].fragments_assembled = conn_stats.fragments_assembled;
                stats->connections[i].current_rtt_ms = conn_stats.current_rtt_ms;
                stats->connections[i].jitter_ms = conn_stats.jitter_ms;
                stats->connections[i].is_active = 1;
            }
        }
        stats->conn_count = max_conn;
    }
    
    // Collect routing statistics
    if (state->routing_table) {
        routing_get_stats(state->routing_table, &stats->routing);
    }
    
    // Collect TUN statistics
    stats->tun.bytes_read = state->tun.bytes_read;
    stats->tun.bytes_written = state->tun.bytes_written;
    stats->tun.packets_read = state->tun.packets_read;
    stats->tun.packets_written = state->tun.packets_written;
    stats->tun.read_errors = state->tun.read_errors;
    stats->tun.write_errors = state->tun.write_errors;
    strncpy(stats->tun.ifname, state->tun.ifname, sizeof(stats->tun.ifname) - 1);
    stats->tun.ifname[sizeof(stats->tun.ifname) - 1] = '\0';
    strncpy(stats->tun.ip_addr, state->tun.ip_addr, sizeof(stats->tun.ip_addr) - 1);
    stats->tun.ip_addr[sizeof(stats->tun.ip_addr) - 1] = '\0';
    
    // Collect system status
    return collect_system_status(cs, &stats->system);
}

// Collect system status
static int collect_system_status(control_socket_t *cs, system_status_t *status) {
    if (!cs) return -1;
    
    memset(status, 0, sizeof(*status));
    
    // Calculate uptime
    struct timeval now;
    gettimeofday(&now, NULL);
    status->uptime_seconds = now.tv_sec - cs->start_time.tv_sec;
    
    // Get memory usage (simplified)
    FILE *meminfo = fopen("/proc/self/statm", "r");
    if (meminfo) {
        unsigned long pages;
        if (fscanf(meminfo, "%lu", &pages) == 1) {
            status->memory_usage_kb = pages * 4; // Assume 4KB pages
        }
        fclose(meminfo);
    }
    
    // CPU usage (placeholder)
    status->cpu_usage_percent = 0;
    
    // Version info
    snprintf(status->version, sizeof(status->version), "utun v0.1");
    
    // Control port
    status->control_port = cs->port;
    
    // Running flag and connection count will be set by caller
    
    return 0;
}

// Get ETCP statistics
static void get_etcp_stats(epkt_t *etcp, etcp_stats_t *etcp_stats) {
    if (!etcp || !etcp_stats) return;
    
    memset(etcp_stats, 0, sizeof(*etcp_stats));
    
    // Use etcp_get_stats if available
    // etcp_get_stats(etcp, &retransmissions, ...);
    
    // Direct access to epkt fields (assuming struct is visible)
    etcp_stats->retransmissions = etcp->retransmissions_count;
    etcp_stats->total_packets_sent = etcp->total_packets_sent;
    etcp_stats->unique_packets_sent = etcp->unique_packets_sent;
    etcp_stats->bytes_sent_total = etcp->bytes_sent_total;
    etcp_stats->bytes_received_total = etcp->bytes_received_total;
    etcp_stats->ack_packets_count = etcp->ack_packets_count;
    etcp_stats->control_packets_count = etcp->control_packets_count;
    etcp_stats->current_rtt = etcp->rtt_avg_10;
    etcp_stats->jitter = etcp->jitter;
    etcp_stats->window_size = etcp->window_size;
    etcp_stats->unacked_bytes = etcp->unacked_bytes;
}

// Get packet normalizer statistics
static void get_pkt_normalizer_stats(pkt_normalizer_pair *pair, 
                                    pkt_normalizer_stats_t *pn_stats) {
    if (!pair || !pn_stats) return;
    
    memset(pn_stats, 0, sizeof(*pn_stats));
    
    if (pair->packer) {
        pn_stats->packer_errors = pkt_normalizer_get_error_count(pair->packer);
    }
    
    if (pair->unpacker) {
        pn_stats->unpacker_errors = pkt_normalizer_get_error_count(pair->unpacker);
        // Fragments assembled count might need to be tracked separately
    }
}