// etcp_loadbalancer.c - Load Balancer Implementation (updated to match etcp_loadbalancer.h)
// Changes:
// - Replaced etcp_loadbalancer_update_after_send with etcp_loadbalancer_send, which now handles link selection, encryption/send, and shaper update.
// - Added loadbalancer_link_ready to notify ETCP_CONN when a link becomes ready (assumes ETCP_CONN has a resume_fn callback; implement in etcp.h/c if needed).
// - Updated constants to match .h (TIMEBASE_NS, SUB_MODULO, etc.).
// - Refined shaper logic for better precision and added normalization.
// - Removed duplicate update_after_send definitions; consolidated into send.
// - Added debug traces for new functions.
// - Assumed ETCP_CONN has a 'resume_fn' field: void (*resume_fn)(struct ETCP_CONN*); // Add to etcp.h struct ETCP_CONN.

#include "etcp_loadbalancer.h"
#include "../lib/debug_config.h"
#include "../lib/u_async.h"
#include <stdlib.h>
#include <stdint.h>  // For UINT64_MAX
#include <time.h>    // For timespec
#include <math.h>    // For rounding
#include "../lib/mem.h"

// Enable comprehensive debug output for loadbalancer module
#define DEBUG_CATEGORY_LOADBALANCER 1
#define ALLOWED_DELTA 5 // 0.5ms allowed gap for time correction

// Forward declarations
static void shaper_timer_cb(void* arg);  // Shaper wait callback

// Select link for transmission
// Algorithm: min inflight_bytes, round-robin among ties
struct ETCP_LINK* etcp_loadbalancer_select_link(struct ETCP_CONN* etcp) {
    if (!etcp || !etcp->links) {
        DEBUG_WARN(DEBUG_CATEGORY_ETCP, "[%s] invalid parameters (etcp=%p, links=%p)", 
                   etcp ? etcp->log_name : "????→????", etcp, etcp ? etcp->links : NULL);
        return NULL;
    }
    
    // First pass: find minimum inflight_bytes among available links
    uint32_t min_inflight = UINT32_MAX;
    int available_count = 0;
    
    struct ETCP_LINK* link = etcp->links;
    while (link) {
        if (!link->initialized || link->link_status != 1) {
            link = link->next;
            continue;
        }
        if (loadbalancer_link_can_send(link) == 0) {
            link = link->next;
            continue;
        }
        
        available_count++;
        if (link->inflight_bytes < min_inflight) {
            min_inflight = link->inflight_bytes;
        }
        link = link->next;
    }
    
    if (available_count == 0) {
        if (etcp->links)
            DEBUG_WARN(DEBUG_CATEGORY_ETCP, "[%s] no suitable link found: inf:%s, tmr:%s", 
                       etcp->log_name, etcp->links->send_blocked_inflight?"wait":"rdy", etcp->links->shaper_timer?"wait":"rdy");
        return NULL;
    }
    
    // Count how many links have the minimum inflight
    int min_count = 0;
    link = etcp->links;
    while (link) {
        if (link->initialized && link->link_status == 1 && 
            loadbalancer_link_can_send(link) && link->inflight_bytes == min_inflight) {
            min_count++;
        }
        link = link->next;
    }
    
    if (min_count == 1) {
        // Only one link with minimum inflight - return it directly
        link = etcp->links;
        while (link) {
            if (link->initialized && link->link_status == 1 && 
                loadbalancer_link_can_send(link) && link->inflight_bytes == min_inflight) {
                return link;
            }
            link = link->next;
        }
    }
    
    // Round-robin among links with minimum inflight
    // Start from the link after last_rr_link (or from the beginning if last_rr_link is NULL)
    struct ETCP_LINK* start = etcp->last_rr_link ? etcp->last_rr_link->next : etcp->links;
    if (!start) start = etcp->links;
    
    link = start;
    do {
        if (link->initialized && link->link_status == 1 && 
            loadbalancer_link_can_send(link) && link->inflight_bytes == min_inflight) {
            etcp->last_rr_link = link;
            return link;
        }
        link = link->next;
        if (!link) link = etcp->links;
    } while (link != start);
    
    // Should not reach here, but return NULL just in case
    return NULL;
}

// New: Send dgram (select link, encrypt/send, update shaper)
void etcp_loadbalancer_send(struct ETCP_DGRAM* dgram) {
    if (!dgram) {
        DEBUG_WARN(DEBUG_CATEGORY_ETCP, "called with NULL dgram");
        return;
    }
    
    struct ETCP_CONN* etcp = dgram->link ? dgram->link->etcp : NULL;
    if (!etcp) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "[%s] no ETCP_CONN associated with dgram", etcp->log_name);
        memory_pool_free(etcp->instance->pkt_pool, dgram);
        return;
    }
    
    // Select link if not already set
    if (!dgram->link) {
        dgram->link = etcp_loadbalancer_select_link(etcp);
        if (!dgram->link) {
            DEBUG_WARN(DEBUG_CATEGORY_ETCP, "[%s] no link available, dropping dgram", etcp->log_name);
            memory_pool_free(etcp->instance->pkt_pool, dgram);
            return;
        }
    }
    
    struct ETCP_LINK* link = dgram->link;
    size_t pkt_size = dgram->data_len + sizeof(uint16_t);  // Include timestamp
    
//    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "[%s] sending dgram on link=%p, size=%zu", etcp->log_name, link, pkt_size);
    
    // Encrypt and send (from etcp_connections.c)
    int send_result = etcp_encrypt_send(dgram);
    if (send_result < 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "[%s] encrypt/send failed (%d)", etcp->log_name, send_result);
        // Don't return here, let the function continue to free dgram at the end
    }
    
    // Update shaper after successful send
    if (link->bandwidth == 0) {
        DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "[%s] unlimited bandwidth, skipping shaper update", etcp->log_name);
        // Don't return here, let the function continue to free dgram at the end
    } else {
        // Time to transmit (ns per byte = 8e6 / bw for Kbits/sec)
        double byte_time_ns = 8000000.0 / (double)link->bandwidth;
        uint64_t tx_time_ns = (uint64_t)(pkt_size * byte_time_ns + 0.5);  // Round nearest
        
        // To sub_nanotime (*10 for 0.1ns)
        uint64_t tx_sub = tx_time_ns * 10;
        
        // Add to load
        link->shaper_sub_nanotime += tx_sub;
        uint64_t carry=link->shaper_sub_nanotime / SUB_MODULO;
        link->shaper_sub_nanotime -= carry * SUB_MODULO;
        link->shaper_load_time_tb += carry;
        
        // Check if exceeds burst -> schedule timer
        uint64_t now_tb=get_time_tb();
        if (link->shaper_load_time_tb >= now_tb + SHAPER_BURST_DELAY_TB) {
            uint64_t wait_tb = link->shaper_load_time_tb - now_tb;
            link->shaper_timer = uasync_set_timeout(link->etcp->instance->ua, wait_tb, link, shaper_timer_cb, "shaper");
            if (link->shaper_timer == NULL) DEBUG_ERROR(DEBUG_CATEGORY_TIMERS, "Filed to allocate timer");
            DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "[%s] scheduled shaper timer (wait_tb=%llu) BW=%d pkt_size=%d", etcp->log_name, (unsigned long long)wait_tb, link->bandwidth, pkt_size);
        }
        
        // Inactivity correction
        if (link->shaper_load_time_tb < now_tb-ALLOWED_DELTA) link->shaper_load_time_tb = now_tb-ALLOWED_DELTA;
    }
//    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "[%s] updated load_tb=%llu, sub=%llu, state=%u", etcp->log_name,
//                (unsigned long long)link->shaper_load_time_tb, (unsigned long long)link->shaper_sub_nanotime, link->shaper_timer==NULL?1:0);
    
    memory_pool_free(etcp->instance->pkt_pool, dgram);
}

int loadbalancer_link_can_send(struct ETCP_LINK* link) {
    if (!link) return 0;
    int can = 1;
    if (link->inflight_bytes >= link->inflight_lim_bytes) {
        link->send_blocked_inflight = 1;
        can = 0;
    } else {
        link->send_blocked_inflight = 0;
    }
    if (link->shaper_timer) {
//        link->send_blocked_bandwidth = 1;
        can = 0;
    } else {
//        link->send_blocked_bandwidth = 0;
    }
    return can;
}

void loadbalancer_link_ready(struct ETCP_LINK* link) {
    if (!link || !link->etcp) {
        DEBUG_WARN(DEBUG_CATEGORY_ETCP, "invalid link (%p)", link);
        return;
    }
    if (loadbalancer_link_can_send(link) == 0) {
        DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "link still blocked");
        return;
    }
    DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "link=%p now ready, notifying ETCP_CONN", link);
    if (link->etcp->link_ready_for_send_fn) {
        link->etcp->link_ready_for_send_fn(link->etcp);
    } else {
        DEBUG_WARN(DEBUG_CATEGORY_ETCP, "no link_ready_for_send_fn set");
    }
}

// Get ETCP link status: 1 = at least one link is up, 0 = all links down or no links
int etcp_loadbalancer_get_link_status(struct ETCP_CONN* etcp) {
    if (!etcp) {
        DEBUG_WARN(DEBUG_CATEGORY_ETCP, "NULL etcp");
        return 0;
    }

    struct ETCP_LINK* link = etcp->links;
    int alive_count = 0;

    while (link) {
        if (link->link_status == 1) {
            alive_count++;
        }
        link = link->next;
    }

    DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "[%s] link status check: %d alive links",
                etcp->log_name ? etcp->log_name : "????→????", alive_count);

    return (alive_count > 0) ? 1 : 0;
}

// Shaper timer callback
static void shaper_timer_cb(void* arg) {
    struct ETCP_LINK* link = (struct ETCP_LINK*)arg;
    if (!link) return;
    
    link->shaper_timer = NULL;
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "link=%p now ready", link);
    
    // Notify loadbalancer that link is ready
    loadbalancer_link_ready(link);
}