// etcp_loadbalancer.c - Load Balancer Implementation (based on etcp_protocol.txt)
#include "etcp_loadbalancer.h"
#include "../lib/debug_config.h"
#include "../lib/u_async.h"
#include <stdlib.h>

// Enable comprehensive debug output for loadbalancer module
#define DEBUG_CATEGORY_LOADBALANCER 1

// Forward declaration
static void init_timeout_cb(void* arg);

// Constants
#define TIMEBASE_NS 100000  // 0.1ms = 100us = 100000ns
#define DELTA_TIME_NS 10000 // Example delta

// Internal
static uint64_t get_current_nanotime() {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    uint64_t nanotime = (uint64_t)ts.tv_sec * 1000000000ULL + ts.tv_nsec;
    DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "get_current_nanotime: tv_sec=%ld, tv_nsec=%ld, result=%llu",
                ts.tv_sec, ts.tv_nsec, (unsigned long long)nanotime);
    return nanotime;
}

// Select link for transmission (per spec)
struct ETCP_LINK* etcp_loadbalancer_select_link(struct ETCP_CONN* etcp) {
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_select_link: etcp=%p", etcp);
    
    if (!etcp || !etcp->links) {
        DEBUG_WARN(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_select_link: invalid parameters (etcp=%p, links=%p)", 
                  etcp, etcp ? etcp->links : NULL);
        return NULL;
    }
    
    struct ETCP_LINK* best = NULL;
    uint64_t min_load_time = UINT64_MAX;
    uint64_t now_ns = get_current_nanotime();
    uint64_t now_tb = now_ns / (TIMEBASE_NS / 10);  // To 0.1ms units
    
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_select_link: current time=%llu tb, scanning %d links", 
                (unsigned long long)now_tb, etcp->links ? 1 : 0);
    
    struct ETCP_LINK* link = etcp->links;
    int link_index = 0;
    while (link) {
        link_index++;
        DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_select_link: evaluating link %d (%p) - initialized=%u, is_server=%u, bandwidth=%u",
                    link_index, link, link->initialized, link->is_server, link->bandwidth);
        
        if (!link->initialized) {
            DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_select_link: link %d not initialized, checking if client connection needed",
                        link_index);
            // Initiate connection if client
            if (!link->is_server && !link->init_timer) {
                DEBUG_INFO(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_select_link: initiating client connection for link %d", link_index);
                // Send INIT (via etcp_encrypt_send with special dgram)
                // Set timer
                link->init_timer = uasync_set_timeout(etcp->instance->ua, link->init_timeout, link, init_timeout_cb);  // Define cb
                DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_select_link: set init timer for link %d", link_index);
            }
            link = link->next;
            continue;
        }
        
        // Update load time if inactive
        if (link->last_activity < now_tb - (DELTA_TIME_NS / (TIMEBASE_NS / 10))) {
            DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_select_link: link %d inactive, updating last_activity from %llu to %llu",
                        link_index, (unsigned long long)link->last_activity, 
                        (unsigned long long)(now_tb - (DELTA_TIME_NS / (TIMEBASE_NS / 10))));
            link->last_activity = now_tb - (DELTA_TIME_NS / (TIMEBASE_NS / 10));
        }
        
        // Check if can send (load time < now)
        DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_select_link: link %d last_activity=%llu, now_tb=%llu",
                    link_index, (unsigned long long)link->last_activity, (unsigned long long)now_tb);
        
        if (link->last_activity < now_tb) {
            DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_select_link: link %d can send (load_time < now)", link_index);
            if (link->last_activity < min_load_time) {
                DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_select_link: link %d is better candidate (activity=%llu < %llu)",
                            link_index, (unsigned long long)link->last_activity, (unsigned long long)min_load_time);
                min_load_time = link->last_activity;
                best = link;
            }
        } else {
            DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_select_link: link %d cannot send yet (load_time >= now)", link_index);
        }
        
        link = link->next;
    }
    
    if (best) {
        DEBUG_INFO(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_select_link: selected link %p (activity=%llu)", 
                  best, (unsigned long long)min_load_time);
    } else {
        DEBUG_WARN(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_select_link: no suitable link found");
    }
    
    return best;
}

// Init timeout callback (placeholder)
static void init_timeout_cb(void* arg) {
    struct ETCP_LINK* link = (struct ETCP_LINK*)arg;
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "init_timeout_cb: link=%p, retry_count=%u, timeout=%u", 
                link, link->init_retry_count, link->init_timeout);
    
    // Resend INIT, increment retry
    link->init_retry_count++;
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "init_timeout_cb: incremented retry_count to %u", link->init_retry_count);
    
    if (link->init_retry_count > 5) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "init_timeout_cb: max retries exceeded (%u > 5), closing link", link->init_retry_count);
        // Fail
        etcp_link_close(link);
        return;
    }
    
    // Resend...
    link->init_timeout *= 2;  // Backoff
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "init_timeout_cb: doubled timeout to %u ms", link->init_timeout);
    link->init_timer = uasync_set_timeout(link->etcp->instance->ua, link->init_timeout, link, init_timeout_cb);
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "init_timeout_cb: rescheduled timer");
}

// Bandwidth limit update (called after send)
void etcp_loadbalancer_update_after_send(struct ETCP_LINK* link, size_t pkt_size) {
    DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_update_after_send: link=%p, pkt_size=%zu", link, pkt_size);
    
    if (!link) {
        DEBUG_WARN(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_update_after_send: called with NULL link");
        return;
    }
    
    if (link->bandwidth == 0) {
        DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_update_after_send: link bandwidth is 0, skipping update");
        return;
    }
    
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_update_after_send: link bandwidth=%u bits/sec, packet size=%zu bytes", 
                link->bandwidth, pkt_size);
    
    // Time to transmit (ns)
    double byte_time_ns = 1000000000.0 / (link->bandwidth * 8.0);  // bits/sec to byte/ns
    uint64_t tx_time_ns = (uint64_t)(pkt_size * byte_time_ns);
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_update_after_send: byte_time_ns=%.2f, tx_time_ns=%llu", 
                byte_time_ns, (unsigned long long)tx_time_ns);
    
    // To timebase (0.1ms units)
    uint64_t tx_time_tb = tx_time_ns / (TIMEBASE_NS / 10);
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_update_after_send: tx_time_tb=%llu", (unsigned long long)tx_time_tb);
    
    // Update last_activity
    uint64_t now_tb = get_current_nanotime() / (TIMEBASE_NS / 10);
    uint64_t old_activity = link->last_activity;
    link->last_activity += tx_time_tb;
    if (link->last_activity < now_tb - (DELTA_TIME_NS / (TIMEBASE_NS / 10))) {
        link->last_activity = now_tb - (DELTA_TIME_NS / (TIMEBASE_NS / 10));
    }
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "etcp_loadbalancer_update_after_send: updated last_activity from %llu to %llu", 
                (unsigned long long)old_activity, (unsigned long long)link->last_activity);
}