// 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
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;
    }
    
    struct ETCP_LINK* best = NULL;
    uint64_t min_load_tb = UINT64_MAX;
    uint64_t now_tb=get_time_tb();
    
    struct ETCP_LINK* link = etcp->links;
    int link_index = 0;
    while (link) {
        link_index++;
        
//        DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_select_link: link %d (%p) - initialized=%u, state=%u, load_tb=%llu, sub=%llu",
//                    link_index, link, link->initialized, link->shaper_timer==NULL?1:0,
//                    (unsigned long long)link->shaper_load_time_tb, (unsigned long long)link->shaper_sub_nanotime);
        
        if (!link->initialized || link->shaper_timer || link->link_status != 1) {// если установлен таймер - значит надо ждать. не рассматриваем этот линк
            link = link->next;
            continue;
        }
        
        if (link->shaper_load_time_tb < now_tb-ALLOWED_DELTA) link->shaper_load_time_tb = now_tb-ALLOWED_DELTA;

        // Check if ready (load < now + burst allowance)
        uint64_t effective_load_ns = link->shaper_load_time_tb * TIMEBASE_NS + (link->shaper_sub_nanotime / 10);  // To ns
        uint64_t now_ns = now_tb * TIMEBASE_NS;
        int64_t delta_t=(int64_t)(now_ns - effective_load_ns);// uint64 может переполняться. чтобы обеспечить правильную цикличность.
        if (link->shaper_load_time_tb < min_load_tb || (link->shaper_load_time_tb == min_load_tb && link->shaper_sub_nanotime < best->shaper_sub_nanotime)) {
            min_load_tb = link->shaper_load_time_tb;
            best = link;
        }
        link = link->next;
    }
    
    if (best) {
//        DEBUG_INFO(DEBUG_CATEGORY_ETCP, "[%s] selected link %p (load_tb=%llu)", 
//                   etcp->log_name, best, (unsigned long long)min_load_tb);
    } else {
        DEBUG_WARN(DEBUG_CATEGORY_ETCP, "[%s] no suitable link found", etcp->log_name);
    }
    
    return best;
}

// New: Send dgram (select link, encrypt/send, update shaper)
void etcp_loadbalancer_send(struct ETCP_DGRAM* dgram) {
    if (!dgram) {
        DEBUG_WARN(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_send: 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);
//        u_free(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);
//            u_free(dgram);  // Assume free; adjust if pooled
            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 = 8e9 / bw for bits/sec)
        double byte_time_ns = 8000000000.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);
            DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "[%s] scheduled shaper timer (wait_tb=%llu)", etcp->log_name, (unsigned long long)wait_tb);
        }
        
        // 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);
//    u_free(dgram);  // Free the dgram in all cases - we own it
}

// New: Notify when link is ready (called from timer or external)
void loadbalancer_link_ready(struct ETCP_LINK* link) {
    if (!link || !link->etcp) {
        DEBUG_WARN(DEBUG_CATEGORY_ETCP, "loadbalancer_link_ready: invalid link (%p)", link);
        return;
    }

    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "loadbalancer_link_ready: link=%p now ready, notifying ETCP_CONN", link);

    // Call ETCP_CONN resume (assumes link_ready_for_send_fn in ETCP_CONN; add to etcp.h: void (*link_ready_for_send_fn)(struct ETCP_CONN*);)
    if (link->etcp->link_ready_for_send_fn) {
        link->etcp->link_ready_for_send_fn(link->etcp);
    } else {
        // Fallback: request next packet (implement etcp_request_pkt in etcp.c if needed)
        // etcp_request_pkt(link->etcp);
        DEBUG_WARN(DEBUG_CATEGORY_ETCP, "loadbalancer_link_ready: no link_ready_for_send_fn set, skipping notify");
    }
}

// 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, "etcp_loadbalancer_get_link_status: 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, "shaper_timer_cb: link=%p now ready", link);
    
    // Notify loadbalancer that link is ready
    loadbalancer_link_ready(link);
}