// routing.c - Centralized routing module for utun
#include "../lib/platform_compat.h"
#include "routing.h"
#include "route_lib.h"
#include "tun_if.h"
#include "packet_dump.h"
#include "etcp.h"
#include "etcp_api.h"
#include "pkt_normalizer.h"
#include "utun_instance.h"
#include "../lib/ll_queue.h"
#include "../lib/debug_config.h"
#include <string.h>
#include <stdlib.h>
#include "../lib/mem.h"

#define IPv4_VERSION 4
#define IPv6_VERSION 6
#define IP_HDR_VERSION_OFFSET 0
#define IP_HDR_DST_ADDR_OFFSET 16
#define IP_HDR_MIN_SIZE 20

// ETCP packet ID for routing data
#define ETCP_ID_DATA 0x00

// Extract destination IP from IPv4 packet
// Returns 0 if not IPv4 or packet too small
static uint32_t extract_dst_ip(uint8_t* data, size_t len) {
    if (!data || len < IP_HDR_MIN_SIZE) {
        return 0;
    }
    
    // Check IP version (first nibble)
    uint8_t version = (data[IP_HDR_VERSION_OFFSET] >> 4) & 0x0F;
    
    if (version == IPv6_VERSION) {
        DEBUG_TRACE(DEBUG_CATEGORY_ROUTING, "IPv6 packet detected, dropping");
        return 0;
    }
    
    if (version != IPv4_VERSION) {
        DEBUG_WARN(DEBUG_CATEGORY_ROUTING, "Unknown IP version: %d", version);
        return 0;
    }
    
    // Extract destination IP (offset 16, 4 bytes, network byte order)
    uint32_t dst_ip;
    memcpy(&dst_ip, data + IP_HDR_DST_ADDR_OFFSET, 4);
    
    return ntohl(dst_ip);  // Convert to host byte order
}

// entry format: <id 1 byte> <ip_packet ...>
static void route_pkt(struct UTUN_INSTANCE* instance, struct ll_entry* entry, uint64_t src_node_id) {
    if (!instance || !entry) {
        DEBUG_ERROR(DEBUG_CATEGORY_ROUTING, "route_pkt: invalid arguments: instance=%p entry=%p entry->len=%zu",
                    (void*)instance, (void*)entry, entry ? entry->len : 0);
        if (entry) {
            queue_entry_free(entry);
            queue_dgram_free(entry);
        }
        return;
    }

    if (!instance->rt) {
        DEBUG_ERROR(DEBUG_CATEGORY_ROUTING, "route_pkt: no route table");
        queue_entry_free(entry);
        queue_dgram_free(entry);
        return;
    }

    if (!instance->tun) {
        DEBUG_ERROR(DEBUG_CATEGORY_ROUTING, "route_pkt: no TUN interface");
        queue_entry_free(entry);
        queue_dgram_free(entry);
        return;
    }

    // Check packet data (skip cmd byte)
    if (!entry->dgram || entry->len < 2) {
        DEBUG_WARN(DEBUG_CATEGORY_ROUTING, "route_pkt: packet too small: len=%zu, expected >=2", entry->len);
        queue_entry_free(entry);
        queue_dgram_free(entry);
        return;
    }

    // Skip cmd byte, extract IP packet (offset 1)
    uint8_t* ip_data = entry->dgram + 1;
    size_t ip_len = entry->len - 1;

    // Check MTU
    if (ip_len > TUN_MAX_PACKET_SIZE) {
        DEBUG_WARN(DEBUG_CATEGORY_ROUTING, "route_pkt: packet too large: len=%zu max=%d, dropping",
                   ip_len, TUN_MAX_PACKET_SIZE);
        queue_entry_free(entry);
        queue_dgram_free(entry);
        return;
    }

    // Extract destination IP
    uint32_t dst_ip = extract_dst_ip(ip_data, ip_len);
    if (dst_ip == 0) {
//        DEBUG_WARN(DEBUG_CATEGORY_ROUTING, "route_pkt: failed to extract dst IP: len=%zu first_byte=0x%02x, dropping",
//                   ip_len, ip_data ? ip_data[0] : 0);
        queue_entry_free(entry);
        queue_dgram_free(entry);
        return;
    }

    // Silently drop broadcast/multicast - these should not be sent over tunnel
    uint8_t first_octet = (dst_ip >> 24) & 0xFF;

    // Multicast: 224.0.0.0 - 239.255.255.255
    if (first_octet >= 224 && first_octet <= 239) {
        queue_entry_free(entry);
        queue_dgram_free(entry);
        return;
    }

    // Broadcast: only 255.255.255.255 (limited broadcast)
    if (dst_ip == 0xFFFFFFFF) {
        queue_entry_free(entry);
        queue_dgram_free(entry);
        return;
    }

    struct in_addr addr;
    addr.s_addr = htonl(dst_ip);
    DEBUG_TRACE(DEBUG_CATEGORY_ROUTING, "route_pkt: dst=%s len=%zu", ip_to_str(&addr, AF_INET).str, ip_len);

    // Lookup route
    struct ROUTE_ENTRY* route = route_lookup(instance->rt, dst_ip);
    if (!route) {
        DEBUG_DEBUG(DEBUG_CATEGORY_ROUTING, "route_pkt: no route to %s, dropping", ip_to_str(&addr, AF_INET).str);
        instance->dropped_packets++;
        queue_entry_free(entry);
        queue_dgram_free(entry);
        return;
    }

    // Determine destination node ID
    uint64_t dst_node_id = instance->node_id; // Default to local node
    struct ETCP_CONN* conn = NULL;
    
    // Check route type: conn_list == NULL means local route
    if (route->conn_list == NULL) {
        // Local route - send to TUN (entry has [cmd=0][IP data], TUN skips cmd byte)
        // dst_node_id is already instance->node_id
    } else {
        // Learned route - use first conn from conn_list
        if (route->conn_list->conninfo_count > 0) {
            conn = route->conn_list->conn_info[0].conn_id;
        }
        if (!conn) {
            DEBUG_WARN(DEBUG_CATEGORY_ROUTING, "route_pkt: route to %s has no next_hop, dropping", ip_to_str(&addr, AF_INET).str);
            instance->dropped_packets++;
            queue_entry_free(entry);
            queue_dgram_free(entry);
            return;
        }
        dst_node_id = conn->peer_node_id;
        
        if (!conn->normalizer) {
            DEBUG_WARN(DEBUG_CATEGORY_ROUTING, "route_pkt: connection for %s has no normalizer, dropping", ip_to_str(&addr, AF_INET).str);
            instance->dropped_packets++;
            queue_entry_free(entry);
            queue_dgram_free(entry);
            return;
        }
    }

    if (debug_should_output(DEBUG_LEVEL_DEBUG, DEBUG_CATEGORY_TRAFFIC)) {
        char* packet_str = dump_ip_packet_to_buffer(ip_data, ip_len);
        DEBUG_DEBUG(DEBUG_CATEGORY_TRAFFIC, "DUMP: %s", packet_str);
    }
    else  DEBUG_INFO(DEBUG_CATEGORY_TRAFFIC, "NODE %016llx -> NODE %016llx", (unsigned long long)src_node_id, (unsigned long long)dst_node_id);

    if (route->conn_list == NULL) {
        // Local route - send to TUN (entry has [cmd=0][IP data], TUN skips cmd byte)
        int put_err = queue_data_put(instance->tun->input_queue, entry, 0);
        if (put_err != 0) {
            DEBUG_WARN(DEBUG_CATEGORY_ROUTING, "route_pkt: failed to put to TUN: dst=%s err=%d",
                       ip_to_str(&addr, AF_INET).str, put_err);
            instance->dropped_packets++;
            queue_entry_free(entry);
            queue_dgram_free(entry);
            return;
        }
        // Entry sent to TUN, don't free here
        instance->routed_packets++;
        DEBUG_TRACE(DEBUG_CATEGORY_ROUTING, "route_pkt: sent %zu bytes to TUN", ip_len);
        return;
    }

    // Send to ETCP
    DEBUG_TRACE(DEBUG_CATEGORY_ROUTING, "route_pkt: sending %zu bytes to ETCP %s", ip_len, conn->log_name);
    int send_err = etcp_send(conn, entry);
    if (send_err != 0) {
        DEBUG_WARN(DEBUG_CATEGORY_ROUTING, "route_pkt: etcp_send failed: dst=%s err=%d",
                   ip_to_str(&addr, AF_INET).str, send_err);
        instance->dropped_packets++;
        queue_entry_free(entry);
        queue_dgram_free(entry);
        return;
    }

    // Entry sent successfully, don't free here
    instance->routed_packets++;
}

// Callback for packets from ETCP (via etcp_bind id=0)
static void routing_pkt_from_etcp_cb(struct ETCP_CONN* conn, struct ll_entry* pkt) {
    if (!conn || !pkt) {
        DEBUG_ERROR(DEBUG_CATEGORY_ROUTING, "routing_pkt_from_etcp_cb: invalid arguments: conn=%p pkt=%p",
                    (void*)conn, (void*)pkt);
        if (pkt) {
            queue_entry_free(pkt);
            queue_dgram_free(pkt);
        }
        return;
    }

    struct UTUN_INSTANCE* instance = conn->instance;
    if (!instance) {
        DEBUG_ERROR(DEBUG_CATEGORY_ROUTING, "routing_pkt_from_etcp_cb: connection has no instance: conn=%p pkt_len=%zu",
                    (void*)conn, pkt->len);
        queue_entry_free(pkt);
        queue_dgram_free(pkt);
        return;
    }

    // Source is the remote node we received the packet from
    uint64_t src_node_id = conn->peer_node_id;
    route_pkt(instance, pkt, src_node_id);
}

// Callback for packets from TUN output queue
static void routing_pkt_from_tun_cb(struct ll_queue* q, void* arg) {
    (void)q;
    struct UTUN_INSTANCE* instance = (struct UTUN_INSTANCE*)arg;
    if (!instance) {
        DEBUG_ERROR(DEBUG_CATEGORY_ROUTING, "routing_pkt_from_tun_cb: invalid instance arg=%p", arg);
        return;
    }

    struct ll_entry* pkt = queue_data_get(instance->tun->output_queue);
    if (!pkt) return;

    // Route packet (unified function handles everything)
    // Source is the local node (TUN interface)
    uint64_t src_node_id = instance->node_id;
    route_pkt(instance, pkt, src_node_id);

    // Resume callback for next packet
    queue_resume_callback(instance->tun->output_queue);
}

// Initialize routing module for instance
int routing_create(struct UTUN_INSTANCE* instance) {
    if (!instance) {
        DEBUG_ERROR(DEBUG_CATEGORY_ROUTING, "routing_create: instance is NULL");
        return -1;
    }
    
    // Create route table
    instance->rt = route_table_create();
    if (!instance->rt) {
        DEBUG_ERROR(DEBUG_CATEGORY_ROUTING, "routing_create: failed to create route table for node %016llx",
                    (unsigned long long)instance->node_id);
        return -1;
    }
    
    // Bind ID=0 for receiving data packets from ETCP
    if (etcp_bind(instance, ETCP_ID_DATA, routing_pkt_from_etcp_cb) != 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_ROUTING, "routing_create: failed to bind ETCP ID=0 for node %016llx",
                    (unsigned long long)instance->node_id);
        route_table_destroy(instance->rt);
        instance->rt = NULL;
        return -1;
    }
    
    DEBUG_INFO(DEBUG_CATEGORY_ROUTING, "Routing module initialized for instance node_id=%016llx",
               (unsigned long long)instance->node_id);
    
    return 0;
}

// Destroy routing module for instance
void routing_destroy(struct UTUN_INSTANCE* instance) {
    if (!instance) return;
    
    DEBUG_INFO(DEBUG_CATEGORY_ROUTING, "Destroying routing module for node_id=%016llx",
               (unsigned long long)instance->node_id);
    
    // Unbind ETCP ID=0
    etcp_unbind(instance, ETCP_ID_DATA);
    
    // Clean up route table if exists
    if (instance->rt) {
        route_table_destroy(instance->rt);
        instance->rt = NULL;
    }
}

// Register ETCP connection with routing (deprecated, kept for compatibility)
void routing_add_conn(struct ETCP_CONN* etcp) {
    (void)etcp;
    // No-op: connections are now handled via routing table entries
    DEBUG_DEBUG(DEBUG_CATEGORY_ROUTING, "routing_add_conn: deprecated, no action taken");
}

// Unregister ETCP connection from routing (deprecated, kept for compatibility)
void routing_del_conn(struct ETCP_CONN* etcp) {
    (void)etcp;
    // No-op: connections are now handled via routing table entries
    DEBUG_DEBUG(DEBUG_CATEGORY_ROUTING, "routing_del_conn: deprecated, no action taken");
}

// Set TUN interface for routing
void routing_set_tun(struct UTUN_INSTANCE* instance) {
    if (!instance) {
        DEBUG_ERROR(DEBUG_CATEGORY_ROUTING, "routing_set_tun: instance is NULL");
        return;
    }
    
    if (!instance->tun) {
        DEBUG_ERROR(DEBUG_CATEGORY_ROUTING, "routing_set_tun: instance %016llx has no TUN",
                    (unsigned long long)instance->node_id);
        return;
    }
    
    // Set callback on TUN output queue (packets from TUN to routing)
    queue_set_callback(instance->tun->output_queue, routing_pkt_from_tun_cb, instance);
    
    DEBUG_INFO(DEBUG_CATEGORY_ROUTING, "TUN interface registered: ifname=%s node_id=%016llx",
               instance->tun->ifname, (unsigned long long)instance->node_id);
}