// utun.c - Main application for utun VPN tunnel
#define _DEFAULT_SOURCE
#define _POSIX_C_SOURCE 200809L
#include "config_parser.h"
#include "connection.h"
#include "tun_if.h"
#include "sc_lib.h"
#include "routing.h"
#include "control_socket.h"
#include "utun_state.h"
#include "u_async.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <errno.h>
#include <signal.h>
#include <poll.h>
#include <getopt.h>
#include <arpa/inet.h>

#define MAX_PACKET_SIZE 2048
#define DEFAULT_CONFIG "utun.conf"
#define DEFAULT_PIDFILE "/var/run/utun.pid"

// Command line arguments
typedef struct {
    char *config_file;
    char *pid_file;
    char *log_file;
    char *tun_ifname;
    char *tun_ip;
    int foreground;
    int help;
} cmd_args_t;

// Global state

// Hex string to binary conversion
static int hex_to_bin(const char *hex, uint8_t *bin, size_t bin_len) {
    if (!hex || !bin) return -1;
    
    size_t hex_len = strlen(hex);
    if (hex_len % 2 != 0 || hex_len / 2 > bin_len) return -1;
    
    for (size_t i = 0; i < hex_len; i += 2) {
        char byte_str[3] = {hex[i], hex[i + 1], '\0'};
        char *endptr;
        long byte = strtol(byte_str, &endptr, 16);
        if (*endptr != '\0') return -1;
        bin[i / 2] = (uint8_t)byte;
    }
    
    return hex_len / 2;
}

// Parse IP:port string
static int parse_addr(const char *addr_str, char *ip, size_t ip_len, uint16_t *port) {
    if (!addr_str || !ip || !port) return -1;
    
    char *colon = strchr(addr_str, ':');
    if (!colon) return -1;
    
    size_t ip_size = colon - addr_str;
    if (ip_size >= ip_len) return -1;
    
    strncpy(ip, addr_str, ip_size);
    ip[ip_size] = '\0';
    
    char *endptr;
    long port_num = strtol(colon + 1, &endptr, 10);
    if (*endptr != '\0' || port_num < 1 || port_num > 65535) return -1;
    
    *port = (uint16_t)port_num;
    return 0;
}

// Parse subnet string
static int parse_subnet(const char *subnet_str, uint32_t *network, uint8_t *prefix_length) {
    if (!subnet_str || !network || !prefix_length) return -1;
    char ip[64];
    int prefix;
    if (sscanf(subnet_str, "%[^/]/%d", ip, &prefix) != 2) return -1;
    if (prefix < 0 || prefix > 32) return -1;
    struct in_addr addr;
    if (inet_pton(AF_INET, ip, &addr) != 1) return -1;
    *network = addr.s_addr;
    *prefix_length = (uint8_t)prefix;
    return 0;
}

// Extract destination IPv4 address from packet
static uint32_t get_dest_ip(const uint8_t *packet, size_t len) {
    if (len < 20) return 0; // Minimum IPv4 header size
    // Check IP version (first nibble)
    uint8_t version = (packet[0] >> 4) & 0x0F;
    if (version != 4) return 0;
    // Destination IP is at offset 16
    uint32_t dest_ip;
    memcpy(&dest_ip, packet + 16, 4);
    return dest_ip;
}

// Initialize connection from configuration
static conn_handle_t* init_connection(uasync_t *ua,
                                      const connection_config_t *conn_cfg,
                                      const global_config_t *global_cfg) {
    if (!conn_cfg || !global_cfg) return NULL;
    
    // Create connection
    conn_handle_t *conn = conn_create(ua);
    if (!conn) {
        fprintf(stderr, "Failed to create connection\n");
        return NULL;
    }
    
    // Convert keys from hex to binary
    uint8_t my_priv_key[32] = {0};
    uint8_t my_pub_key[64] = {0};
    uint8_t peer_pub_key[64] = {0};
    
    if (strlen(global_cfg->my_private_key_hex) > 0) {
        if (hex_to_bin(global_cfg->my_private_key_hex, my_priv_key, sizeof(my_priv_key)) < 0) {
            fprintf(stderr, "Invalid private key format\n");
            conn_destroy(conn);
            return NULL;
        }
    }
    
    if (strlen(global_cfg->my_public_key_hex) > 0) {
        if (hex_to_bin(global_cfg->my_public_key_hex, my_pub_key, sizeof(my_pub_key)) < 0) {
            fprintf(stderr, "Invalid public key format\n");
            conn_destroy(conn);
            return NULL;
        }
    }
    
    if (strlen(conn_cfg->peer_public_key_hex) > 0) {
        if (hex_to_bin(conn_cfg->peer_public_key_hex, peer_pub_key, sizeof(peer_pub_key)) < 0) {
            fprintf(stderr, "Invalid peer public key format\n");
            conn_destroy(conn);
            return NULL;
        }
    }
    
    // Set keys
    if (conn_set_keys(conn, my_pub_key, my_priv_key, peer_pub_key) < 0) {
        fprintf(stderr, "Failed to set keys\n");
        conn_destroy(conn);
        return NULL;
    }
    
    // Parse addresses
    char local_ip[64] = "";
    uint16_t local_port = 0;
    char remote_ip[64] = "";
    uint16_t remote_port = 0;
    
    if (strlen(conn_cfg->local_addr) > 0) {
        if (parse_addr(conn_cfg->local_addr, local_ip, sizeof(local_ip), &local_port) < 0) {
            fprintf(stderr, "Invalid local address format: %s\n", conn_cfg->local_addr);
            conn_destroy(conn);
            return NULL;
        }
    }
    
    if (conn_cfg->mode == CONFIG_MODE_CLIENT && strlen(conn_cfg->remote_addr) > 0) {
        if (parse_addr(conn_cfg->remote_addr, remote_ip, sizeof(remote_ip), &remote_port) < 0) {
            fprintf(stderr, "Invalid remote address format: %s\n", conn_cfg->remote_addr);
            conn_destroy(conn);
            return NULL;
        }
        
        // Apply net_debug mode: redirect traffic to emulator port (+10000)
        if (global_cfg->net_debug && remote_port > 0) {
            uint16_t original_port = remote_port;
            remote_port += 10000;
            if (remote_port < original_port) { // Overflow check
                remote_port = 65535;
            }
            printf("[NET_DEBUG] Redirecting connection %s:%u -> %s:%u (emulator port %u)\n",
                   conn_cfg->name, original_port, remote_ip, remote_port - 10000, remote_port);
        }
    }
    
    // Connect
    conn_mode_t mode = (conn_cfg->mode == CONFIG_MODE_CLIENT) ? CONN_MODE_CLIENT : CONN_MODE_SERVER;
    
    if (conn_connect(conn, 
                     strlen(local_ip) > 0 ? local_ip : NULL,
                     local_port,
                     strlen(remote_ip) > 0 ? remote_ip : NULL,
                     remote_port,
                     mode) < 0) {
        fprintf(stderr, "Failed to connect\n");
        conn_destroy(conn);
        return NULL;
    }
    
    return conn;
}

// Callback for received data from connection
static void connection_recv_callback(conn_handle_t* conn, 
                                     const uint8_t* data, 
                                     size_t len,
                                     void* user_data) {
    (void)conn; // unused
    utun_state_t *state = (utun_state_t*)user_data;
    if (!state || !data || len == 0) return;
    
    // Write to TUN device
    ssize_t written = tun_write(state->tun.fd, data, len);
    if (written > 0) {
        state->tun.bytes_written += written;
        state->tun.packets_written++;
    } else if (written < 0) {
        state->tun.write_errors++;
        fprintf(stderr, "Failed to write to TUN device: %s\n", strerror(errno));
    } else if (written != (ssize_t)len) {
        fprintf(stderr, "Partial write to TUN device: %zd/%zu\n", written, len);
    }
}

// Initialize all connections
static int init_connections(utun_state_t *state) {
    if (!state || !state->config) return -1;
    
    state->connection_count = state->config->connection_count;
    state->connections = calloc(state->connection_count, sizeof(conn_handle_t*));
    if (!state->connections) return -1;
    
    for (int i = 0; i < state->connection_count; i++) {
        state->connections[i] = init_connection(state->ua, 
                                                 &state->config->connections[i], 
                                                 &state->config->global);
        if (!state->connections[i]) {
            fprintf(stderr, "Failed to initialize connection %d\n", i);
            // Cleanup already created connections
            for (int j = 0; j < i; j++) {
                conn_destroy(state->connections[j]);
            }
            free(state->connections);
            state->connections = NULL;
            state->connection_count = 0;
            return -1;
        }
        
        // Set receive callback
        conn_set_recv_callback(state->connections[i], 
                               connection_recv_callback, 
                               state);
    }
    
    return 0;
}

// Add default route via first connection
static int add_default_route(utun_state_t *state) {
    if (!state || !state->routing_table || state->connection_count == 0) return -1;
    
    route_entry_t default_route = {0};
    default_route.network = 0; // 0.0.0.0
    default_route.prefix_length = 0;
    default_route.next_hop_ip = 0; // Unknown
    default_route.next_hop = state->connections[0];
    default_route.type = ROUTE_TYPE_STATIC;
    default_route.flags = ROUTE_FLAG_ACTIVE;
    default_route.created_time = 0; // Not used
    default_route.last_update = 0;
    default_route.last_used = 0;
    
    if (!routing_table_insert(state->routing_table, &default_route)) {
        fprintf(stderr, "Failed to add default route\n");
        return -1;
    }
    
    printf("Added default route via first connection\n");
    return 0;
}

// Parse command line arguments
static void parse_args(int argc, char *argv[], cmd_args_t *args) {
    memset(args, 0, sizeof(*args));
    args->config_file = DEFAULT_CONFIG;
    args->pid_file = DEFAULT_PIDFILE;
    args->foreground = 0;
    args->help = 0;
    
    static struct option long_options[] = {
        {"config", required_argument, 0, 'c'},
        {"pidfile", required_argument, 0, 'p'},
        {"log", required_argument, 0, 'l'},
        {"tun", required_argument, 0, 't'},
        {"tun-ip", required_argument, 0, 'i'},
        {"foreground", no_argument, 0, 'f'},
        {"help", no_argument, 0, 'h'},
        {0, 0, 0, 0}
    };
    
    int opt;
    int option_index = 0;
    
    while ((opt = getopt_long(argc, argv, "c:p:l:t:i:fh", 
                              long_options, &option_index)) != -1) {
        switch (opt) {
            case 'c':
                args->config_file = optarg;
                break;
            case 'p':
                args->pid_file = optarg;
                break;
            case 'l':
                args->log_file = optarg;
                break;
            case 't':
                args->tun_ifname = optarg;
                break;
            case 'i':
                args->tun_ip = optarg;
                break;
            case 'f':
                args->foreground = 1;
                break;
            case 'h':
                args->help = 1;
                break;
            default:
                fprintf(stderr, "Unknown option: %c\n", opt);
                exit(1);
        }
    }
}

// Print usage
static void print_usage(const char *progname) {
    printf("Usage: %s [OPTIONS]\n", progname);
    printf("Secure VPN tunnel over UDP with TUN interface\n\n");
    printf("Options:\n");
    printf("  -c, --config FILE    Configuration file (default: %s)\n", DEFAULT_CONFIG);
    printf("  -p, --pidfile FILE   PID file (default: %s)\n", DEFAULT_PIDFILE);
    printf("  -l, --log FILE       Log file (default: stderr)\n");
    printf("  -t, --tun IFNAME     TUN interface name (e.g., tun12)\n");
    printf("  -i, --tun-ip ADDR    TUN interface IP address (e.g., 10.0.0.1/24)\n");
    printf("  -f, --foreground     Run in foreground (don't daemonize)\n");
    printf("  -h, --help           Show this help\n");
    printf("\nExamples:\n");
    printf("  %s -c myconfig.conf -t tun0 -i 10.0.0.1/24\n", progname);
    printf("  %s --config server.conf --pidfile /var/run/utun.pid\n", progname);
}

// Write PID file
static int write_pidfile(const char *pidfile) {
    if (!pidfile) return -1;
    
    FILE *fp = fopen(pidfile, "w");
    if (!fp) {
        perror("fopen pidfile");
        return -1;
    }
    
    fprintf(fp, "%d\n", getpid());
    fclose(fp);
    return 0;
}

// Remove PID file
static void remove_pidfile(const char *pidfile) {
    if (pidfile) {
        unlink(pidfile);
    }
}

// Daemonize process
static int daemonize(void) {
    pid_t pid = fork();
    if (pid < 0) {
        perror("fork");
        return -1;
    }
    
    if (pid > 0) {
        // Parent exits
        exit(0);
    }
    
    // Child becomes session leader
    if (setsid() < 0) {
        perror("setsid");
        return -1;
    }
    
    // Close standard file descriptors
    close(STDIN_FILENO);
    close(STDOUT_FILENO);
    close(STDERR_FILENO);
    
    // Redirect to /dev/null
    int fd = open("/dev/null", O_RDWR);
    if (fd >= 0) {
        dup2(fd, STDIN_FILENO);
        dup2(fd, STDOUT_FILENO);
        dup2(fd, STDERR_FILENO);
        if (fd > 2) close(fd);
    }
    
    return 0;
}

// Open log file
static FILE* open_logfile(const char *logfile) {
    if (!logfile) return stderr;
    
    FILE *fp = fopen(logfile, "a");
    if (!fp) {
        perror("fopen logfile");
        return stderr;
    }
    
    // Set line buffering
    setlinebuf(fp);
    return fp;
}

// Cleanup function
static void cleanup(utun_state_t *state, const char *pidfile) {
    if (!state) return;
    
    state->running = 0;
    
    // Destroy connections
    if (state->connections) {
        for (int i = 0; i < state->connection_count; i++) {
            if (state->connections[i]) {
                conn_destroy(state->connections[i]);
            }
        }
        free(state->connections);
        state->connections = NULL;
    }

    // Destroy routing table
    if (state->routing_table) {
        routing_table_destroy(state->routing_table);
        state->routing_table = NULL;
    }
    
    // Destroy control socket
    if (state->control_socket) {
        control_socket_destroy(state->control_socket);
        state->control_socket = NULL;
    }

    // Destroy uasync instance
    if (state->ua) {
        uasync_destroy(state->ua);
        state->ua = NULL;
    }
    
    // Close TUN device
    tun_close(&state->tun);
    
    // Free config
    if (state->config) {
        free_config(state->config);
        state->config = NULL;
    }
    
    // Close log file if not stderr
    if (state->log_fp && state->log_fp != stderr) {
        fclose(state->log_fp);
        state->log_fp = NULL;
    }
    
    // Remove PID file
    remove_pidfile(pidfile);
}

// Signal handler
static volatile sig_atomic_t got_signal = 0;
static void signal_handler(int sig) {
    (void)sig; // unused
    got_signal = 1;
}

// Setup signal handlers
static void setup_signals(void) {
    struct sigaction sa;
    sa.sa_handler = signal_handler;
    sa.sa_flags = 0;
    sigemptyset(&sa.sa_mask);
    
    sigaction(SIGINT, &sa, NULL);
    sigaction(SIGTERM, &sa, NULL);
    sigaction(SIGHUP, &sa, NULL);
}

// Main event loop
static int event_loop(utun_state_t *state) {
    // Determine number of file descriptors to poll
    int num_fds = 1; // TUN device always
    int ctrl_fd_idx = -1;
    
    if (state->control_socket) {
        num_fds = 2;
        ctrl_fd_idx = 1;
    }
    
    struct pollfd fds[2];
    uint8_t buffer[MAX_PACKET_SIZE];
    
    // Setup poll for TUN device
    fds[0].fd = state->tun.fd;
    fds[0].events = POLLIN;
    
    // Setup poll for control socket if exists
    if (state->control_socket) {
        fds[ctrl_fd_idx].fd = control_socket_get_fd(state->control_socket);
        fds[ctrl_fd_idx].events = POLLIN;
    }
    
    while (state->running && !got_signal) {
        // Process async events (timers, sockets)
        if (state->ua) {
            uasync_poll(state->ua, 0);
        }
        // Poll all file descriptors with timeout
        int ret = poll(fds, num_fds, 100); // 100ms timeout
        if (ret < 0) {
            if (errno == EINTR) continue;
            perror("poll");
            break;
        }
        
        if (ret == 0) {
            // Timeout - continue
            continue;
        }
        
        // Check control socket first
        if (state->control_socket && (fds[ctrl_fd_idx].revents & POLLIN)) {
            if (control_socket_process(state->control_socket, state) < 0) {
                fprintf(stderr, "Error processing control socket request\n");
            }
        }
        
        // Check TUN device
        if (fds[0].revents & POLLIN) {
            // Read from TUN device
            ssize_t nread = tun_read(state->tun.fd, buffer, sizeof(buffer));
            if (nread < 0) {
                if (errno == EINTR) continue;
                perror("read from TUN");
                state->tun.read_errors++;
                break;
            }
            
            if (nread > 0) {
                state->tun.bytes_read += nread;
                state->tun.packets_read++;
                // Route packet based on destination IP
                uint32_t dest_ip = get_dest_ip(buffer, nread);
                route_entry_t route;
                if (dest_ip != 0 && routing_table_lookup(state->routing_table, dest_ip, &route)) {
                    // Found route, send to next hop connection
                    if (route.next_hop) {
                        if (conn_send(route.next_hop, buffer, nread) < 0) {
                            fprintf(stderr, "Failed to send packet via route\n");
                        }
                    } else {
                        // Local route, no forwarding needed
                    }
                } else {
                    // No route found, drop packet
                    char ip_str[16];
                    ip_to_string(dest_ip, ip_str);
                    fprintf(stderr, "No route for destination IP %s\n", ip_str);
                }
            }
        }
    }
    
    return 0;
}

// Main function
int main(int argc, char *argv[]) {
    cmd_args_t args;
    parse_args(argc, argv, &args);
    
    if (args.help) {
        print_usage(argv[0]);
        return 0;
    }
    
    // Parse configuration
    utun_config_t *config = parse_config(args.config_file);
    if (!config) {
        fprintf(stderr, "Failed to parse configuration file: %s\n", args.config_file);
        return 1;
    }
    
    // Print config for debugging
    if (args.foreground) {
        print_config(config);
    }
    
    // Log net_debug mode
    if (config->global.net_debug) {
        printf("[NET_DEBUG] Network debug mode ENABLED\n");
        printf("[NET_DEBUG] Outgoing client connections will use port +10000\n");
    }
    
    // Initialize state
    utun_state_t state = {0};
    state.config = config;
    state.running = 1;
    state.tun.fd = -1;
    state.routing_table = routing_table_create();
    if (!state.routing_table) {
        fprintf(stderr, "Failed to create routing table\n");
        cleanup(&state, args.pid_file);
        return 1;
    }

    state.ua = uasync_create();
    if (!state.ua) {
        fprintf(stderr, "Failed to create uasync instance\n");
        cleanup(&state, args.pid_file);
        return 1;
    }
    uasync_init_instance(state.ua);
    
    // Setup TUN configuration from command line or config
    if (args.tun_ifname) {
        strncpy(state.tun.ifname, args.tun_ifname, sizeof(state.tun.ifname) - 1);
    }
    if (args.tun_ip) {
        strncpy(state.tun.ip_addr, args.tun_ip, sizeof(state.tun.ip_addr) - 1);
    }
    
    // If not set by command line, try to get from first connection config
    if (state.tun.ifname[0] == '\0' && config->connection_count > 0) {
        if (config->connections[0].tun_ifname[0] != '\0') {
            strncpy(state.tun.ifname, config->connections[0].tun_ifname, sizeof(state.tun.ifname) - 1);
        }
    }
    if (state.tun.ip_addr[0] == '\0' && config->connection_count > 0) {
        if (config->connections[0].tun_ip[0] != '\0') {
            strncpy(state.tun.ip_addr, config->connections[0].tun_ip, sizeof(state.tun.ip_addr) - 1);
        }
    }
    
    state.tun.mtu = 1500; // Default MTU

    // Load allowed subnets into routing table
    for (int i = 0; i < config->allowed_subnet_count; i++) {
        uint32_t network;
        uint8_t prefix_len;
        if (parse_subnet(config->allowed_subnets[i].subnet, &network, &prefix_len) == 0) {
            // Add as dynamic subnet (allowed for validation)
            if (!routing_add_dynamic_subnet(state.routing_table, network, prefix_len)) {
                fprintf(stderr, "Failed to add allowed subnet: %s\n", config->allowed_subnets[i].subnet);
            } else {
                printf("Added allowed subnet: %s\n", config->allowed_subnets[i].subnet);
            }
        } else {
            fprintf(stderr, "Invalid subnet format: %s\n", config->allowed_subnets[i].subnet);
        }
    }
    
    // Add TUN interface subnet as local subnet
    if (state.tun.ip_addr[0] != '\0') {
        uint32_t network;
        uint8_t prefix_len;
        if (parse_subnet(state.tun.ip_addr, &network, &prefix_len) == 0) {
            if (!routing_add_local_subnet(state.routing_table, network, prefix_len)) {
                fprintf(stderr, "Failed to add TUN subnet as local: %s\n", state.tun.ip_addr);
            } else {
                printf("Added TUN subnet as local: %s\n", state.tun.ip_addr);
            }
        } else {
            fprintf(stderr, "Invalid TUN IP format: %s\n", state.tun.ip_addr);
        }
    }
    
    // Initialize control socket if configured
    if (config->global.control_port != 0) {
        const char *control_ip = config->global.control_ip[0] != '\0' ? config->global.control_ip : NULL;
        state.control_socket = control_socket_create(control_ip, config->global.control_port);
        if (!state.control_socket) {
            fprintf(stderr, "Failed to create control socket on %s:%u\n",
                   control_ip ? control_ip : "0.0.0.0", config->global.control_port);
            cleanup(&state, args.pid_file);
            return 1;
        }
        control_socket_set_state(state.control_socket, &state);
        printf("Control socket listening on %s:%u\n",
               control_ip ? control_ip : "0.0.0.0", config->global.control_port);
    }
    
    // Open log file
    state.log_fp = open_logfile(args.log_file);
    
    // Daemonize if not in foreground
    if (!args.foreground) {
        if (daemonize() < 0) {
            fprintf(stderr, "Failed to daemonize\n");
            cleanup(&state, args.pid_file);
            return 1;
        }
    }
    
    // Write PID file
    if (write_pidfile(args.pid_file) < 0) {
        cleanup(&state, args.pid_file);
        return 1;
    }
    
    // Setup signal handlers
    setup_signals();
    
    // Create TUN device
    if (tun_create(&state.tun) < 0) {
        fprintf(stderr, "Failed to create TUN device\n");
        cleanup(&state, args.pid_file);
        return 1;
    }
    
    fprintf(state.log_fp, "Created TUN device: %s\n", state.tun.ifname);
    
    // Initialize connections
    if (init_connections(&state) < 0) {
        fprintf(stderr, "Failed to initialize connections\n");
        cleanup(&state, args.pid_file);
        return 1;
    }
    
    fprintf(state.log_fp, "Initialized %d connection(s)\n", state.connection_count);
    
    // Add default route via first connection
    add_default_route(&state);

    // Run event loop
    fprintf(state.log_fp, "Starting event loop...\n");
    event_loop(&state);
    
    // Cleanup
    fprintf(state.log_fp, "Shutting down...\n");
    cleanup(&state, args.pid_file);
    
    return 0;
}