// etcp.c - ETCP Protocol Implementation
#include "../lib/u_async.h"
#include "utun_instance.h"
#include "etcp.h"
#include "../lib/debug_config.h"
#include "../lib/ll_queue.h"
#include "crc32.h"
#include "secure_channel.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <sys/time.h>
#include <time.h>

// Service packet headers (from protocol, but reset/init handled elsewhere)
#define ETCP_ACK_HEADER 0x01
#define ETCP_RETRANS_HEADER_BASE 0x10  // 0x10 to 0x2F for retrans requests
#define ETCP_PAYLOAD_HEADER 0x00
#define ETCP_INIT_RESPONSE 0x03

// Forward declarations of internal functions (adapted from etcp_master)
static void retransmit_check(struct ETCP_CONN* etcp);
static void update_metrics(struct ETCP_CONN* etcp, uint16_t rtt);

static uint16_t timestamp_diff(uint16_t t1, uint16_t t2);
static int id_compare(uint16_t id1, uint16_t id2);
static void retransmit_packet(struct ETCP_CONN* etcp, uint16_t id);

// Timer callbacks
static void tx_timer_callback(void* arg);
static void retransmit_timer_callback(void* arg);

// Internal queue callbacks
static void tx_queue_callback(struct ll_queue* q, struct ll_entry* entry, void* arg);
static void schedule_ack_timer(struct ETCP_CONN* etcp);
static void request_retransmission_for_gaps(struct ETCP_CONN* etcp);

static void queue_clear(struct ll_queue* q) {
    struct ll_entry* entry;
    while ((entry = queue_entry_get(q)) != NULL) {
        queue_entry_free(entry);
    }
}

static void etcp_update_window(struct ETCP_CONN* etcp) {
    etcp->window_size = 65536; // Initial window size
    etcp->retrans_timer_period = 20; // Default retransmit period
}

// Reset connection (adapted from etcp_reset in master, without packet sending)
void etcp_conn_reset(struct ETCP_CONN* etcp) {
    if (!etcp) return;
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "etcp_conn_reset: resetting instance");
    
    // Cancel timers
    if (etcp->next_tx_timer) {
        uasync_cancel_timeout(etcp->instance->ua, etcp->next_tx_timer);
        etcp->next_tx_timer = NULL;
    }
    if (etcp->retransmit_timer) {
        uasync_cancel_timeout(etcp->instance->ua, etcp->retransmit_timer);
        etcp->retransmit_timer = NULL;
    }
    
    // Clear queues (but keep them)
    queue_clear(etcp->input_queue);
    queue_clear(etcp->output_queue);
    
    // Clear lists
    rx_packet_t* rx = etcp->rx_list;
    while (rx) {
        rx_packet_t* next = rx->next;
        if (rx->data) free(rx->data);
        free(rx);
        rx = next;
    }
    etcp->rx_list = NULL;
    
    sent_packet_t* sent = etcp->sent_list;
    while (sent) {
        sent_packet_t* next = sent->next;
        if (sent->data) free(sent->data);
        free(sent);
        sent = next;
    }
    etcp->sent_list = NULL;
    
    // Reset metrics and stats
    etcp->rtt_last = 0;
    etcp->rtt_avg_10 = 0;
    etcp->rtt_avg_100 = 0;
    etcp->jitter = 0;
    etcp->bytes_sent_total = 0;
    etcp->retransmissions_count = 0;
    etcp->ack_packets_count = 0;
    etcp->control_packets_count = 0;
    etcp->total_packets_sent = 0;
    etcp->unique_packets_sent = 0;
    etcp->bytes_received_total = 0;
    
    // Reset IDs
    etcp->next_tx_id = 1;
    etcp->last_sent_id = 0;
    etcp->last_rx_id = 0;
    etcp->last_delivered_id = 0;
    
    // Reset history
    etcp->rtt_history_idx = 0;
    etcp->rtt_history_count = 0;
    
    // Reset pending
    etcp->pending_ack_count = 0;
    etcp->pending_retransmit_count = 0;
    
    // Reset window
    etcp->unacked_bytes = 0;
    etcp->last_acked_id = 0;
    etcp->last_rx_ack_id = 0;
    etcp->next_retrans_time = 0;
    etcp->window_blocked = 0;
    
    // Reset forward progress
    etcp->oldest_missing_id = 0;
    etcp->missing_since_time = 0;
    
    etcp_update_window(etcp);
}

// Creating ETCP instance
// после создания надо добавить peer bublic key.
struct ETCP_CONN* etcp_connection_create(struct UTUN_INSTANCE* instance) {
    if (!instance) return NULL;
    
    struct ETCP_CONN* etcp = calloc(1, sizeof(struct ETCP_CONN));
    if (!etcp) return NULL;
    
    etcp->mtu = 1500; // Default MTU
    
    etcp->instance = instance;

    // Initialize crypto context
    if (sc_init_ctx(&etcp->crypto_ctx, &etcp->instance->my_keys) != SC_OK) {
        DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, "etcp_connection_create: failed to initialize crypto context for node %llu", (unsigned long long)instance->node_id);
        free(etcp);
        return NULL;
    }

    // Initialize queues (tx_queue is input_queue)
    etcp->input_queue = queue_new(instance->ua,NULL);
    if (!etcp->input_queue) {
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "etcp_connection_create: failed to create input queue for node %llu", (unsigned long long)instance->node_id);
        free(etcp);
        return NULL;
    }
    
    etcp->output_queue = queue_new(instance->ua,NULL);
    if (!etcp->output_queue) {
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "etcp_connection_create: failed to create output queue for node %llu", (unsigned long long)instance->node_id);
        queue_free(etcp->input_queue);
        free(etcp);
        return NULL;
    }
    
    // Set callback for input_queue (tx_queue)
    queue_set_callback(etcp->input_queue, tx_queue_callback, etcp);
    
    // Initialize state (from master)
    etcp->bandwidth = 10000;   // Default: 10000 bytes per timebase (0.1us)
    etcp->last_sent_timestamp = get_current_timestamp();
    etcp->bytes_allowed = 0;
    
    // Reset all other fields to initial state
    etcp_conn_reset(etcp);

    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "etcp_connection_create: created instance with node_id=%llu, mtu=%d", 
               (unsigned long long)etcp->instance->node_id, etcp->mtu);

    etcp->next=instance->connections;
    instance->connections=etcp;

    instance->connections_count++;
    
    return etcp;
}

// Destroying ETCP instance
void etcp_connection_close(struct ETCP_CONN* etcp) {
    if (!etcp) return;
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "etcp_connection_close: destroying instance");
    
    // Cancel timers
    if (etcp->next_tx_timer) {
        uasync_cancel_timeout(etcp->instance->ua, etcp->next_tx_timer);
        etcp->next_tx_timer = NULL;
    }
    if (etcp->retransmit_timer) {
        uasync_cancel_timeout(etcp->instance->ua, etcp->retransmit_timer);
        etcp->retransmit_timer = NULL;
    }
    
    // Free queues
    if (etcp->input_queue) queue_free(etcp->input_queue);
    if (etcp->output_queue) queue_free(etcp->output_queue);
    
    // Free rx_list
    rx_packet_t* rx = etcp->rx_list;
    while (rx) {
        rx_packet_t* next = rx->next;
        if (rx->data) free(rx->data);
        free(rx);
        rx = next;
    }
    
    // Free sent_list
    sent_packet_t* sent = etcp->sent_list;
    while (sent) {
        sent_packet_t* next = sent->next;
        if (sent->data) free(sent->data);
        free(sent);
        sent = next;
    }
    
    free(etcp);
}

// Process incoming packet (partial, truncated in original)
void etcp_conn_input(struct ETCP_DGRAM* pkt) {
    if (!pkt) return;
    struct ETCP_CONN* etcp=pkt->link->etcp;
    if (!etcp) return;
    
    uint8_t* data = pkt->data;
    size_t len = pkt->data_len;
    
    if (len < 4) {
        DEBUG_WARN(DEBUG_CATEGORY_ETCP, "etcp_conn_input: packet too short (%zu bytes) from node %llu", len, (unsigned long long)etcp->instance->node_id);
        return; // Min header
    }
    
    uint16_t id = (data[0] << 8) | data[1];
    uint16_t timestamp = (data[2] << 8) | data[3];
    
    // Check for INIT_RESPONSE packet (special control packet)
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "[ETCP DEBUG] etcp_conn_input: id=%u, len=%zu, first bytes: %02x %02x %02x %02x %02x",
                id, len, data[0], data[1], data[2], data[3], data[4]);
    
    data += 4; len -= 4;  // Skip header for regular packet processing
    
    int has_payload = 0;
    uint8_t* payload_data = NULL;
    size_t payload_len = 0;
    int is_duplicate = 0;
    
    while (len > 0) {
        uint8_t section_header = *data++;
        len--;
        
        if (section_header == ETCP_ACK_HEADER) {
            // ACK section
            uint8_t count = *data++;
            len--;
            for (uint8_t i = 0; i < count; i++) {
                uint16_t ack_id = (data[0] << 8) | data[1];
                uint16_t ack_ts = (data[2] << 8) | data[3];
                data += 4; len -= 4;
                
                // Process ACK: remove from sent_list, update window
                sent_packet_t** ptr = &etcp->sent_list;
                while (*ptr) {
                    if ((*ptr)->id == ack_id) {
                        sent_packet_t* to_free = *ptr;
                        *ptr = to_free->next;
                        etcp->unacked_bytes -= to_free->payload_len;
                        update_metrics(etcp, timestamp_diff(get_current_timestamp(), ack_ts));
                        free(to_free->data);
                        free(to_free);
                        break;
                    }
                    ptr = &(*ptr)->next;
                }
            }
            
            // last_delivered and last_rx
            uint16_t last_delivered = (data[0] << 8) | data[1];
            uint16_t last_rx = (data[2] << 8) | data[3];
            data += 4; len -= 4;
            
            etcp->last_rx_ack_id = last_rx;
            
        } else if ((section_header & 0xF0) == ETCP_RETRANS_HEADER_BASE) {
            // Retrans request
            uint8_t count = (section_header & 0x0F) + 1;
            for (uint8_t i = 0; i < count; i++) {
                uint16_t retrans_id = (data[0] << 8) | data[1];
                data += 2;
                len -= 2;
                
                retransmit_packet(etcp, retrans_id);
            }
            
            // last_delivered and last_rx
            uint16_t last_delivered = (data[0] << 8) | data[1];
            uint16_t last_rx = (data[2] << 8) | data[3];
            data += 4; len -= 4;
            
        } else if (section_header == ETCP_PAYLOAD_HEADER) {
            // Payload section
            has_payload = 1;
            payload_data = data;
            payload_len = len;
            break;  // Payload is last
        }
    }
    
    // Check for duplicate
    rx_packet_t* rx = etcp->rx_list;
    while (rx) {
        if (rx->id == id) {
            is_duplicate = 1;
            break;
        }
        rx = rx->next;
    }
    
    if (is_duplicate) {
        // Add to pending ACKs
        if (etcp->pending_ack_count < 32) {
            etcp->pending_ack_ids[etcp->pending_ack_count] = id;
            etcp->pending_ack_timestamps[etcp->pending_ack_count] = timestamp;
            etcp->pending_ack_count++;
        }
        schedule_ack_timer(etcp);
        return;
    }
    
    // Insert into rx_list (sorted)
    rx_packet_t* new_rx = calloc(1, sizeof(rx_packet_t));
    if (!new_rx) {
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "etcp_conn_input: failed to allocate memory for rx packet from node %llu", (unsigned long long)etcp->instance->node_id);
        return;
    }
    
    new_rx->id = id;
    new_rx->timestamp = timestamp;
    new_rx->has_payload = has_payload;
    if (has_payload) {
        new_rx->data = malloc(payload_len);
        if (!new_rx->data) {
            DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "etcp_conn_input: failed to allocate %zu bytes for payload data from node %llu", payload_len, (unsigned long long)etcp->instance->node_id);
            free(new_rx);
            return;
        }
        memcpy(new_rx->data, payload_data, payload_len);
        new_rx->data_len = payload_len;
        etcp->bytes_received_total += payload_len;
    }
    
    // Insert sorted
    rx_packet_t** ptr = &etcp->rx_list;
    while (*ptr && id_compare((*ptr)->id, id) < 0) {
        ptr = &(*ptr)->next;
    }
    new_rx->next = *ptr;
    *ptr = new_rx;
    
    // Update last_rx_id
    if (id_compare(id, etcp->last_rx_id) > 0) {
        etcp->last_rx_id = id;
    }
    
    // Add to pending ACKs
    if (etcp->pending_ack_count < 32) {
        etcp->pending_ack_ids[etcp->pending_ack_count] = id;
        etcp->pending_ack_timestamps[etcp->pending_ack_count] = timestamp;
        etcp->pending_ack_count++;
    }
    
    // Request retrans for gaps and deliver contiguous
    request_retransmission_for_gaps(etcp);
    schedule_ack_timer(etcp);
}


// Getting statistics (extended with master stats)
void etcp_get_stats(struct ETCP_CONN* etcp, size_t* packets_sent, size_t* packets_recv,
                    size_t* pool_allocs, size_t* pool_reuse) {
    if (!etcp) return;
    
    if (packets_sent) *packets_sent = etcp->total_packets_sent;
    if (packets_recv) *packets_recv = etcp->bytes_received_total / 100; // Approximate
    
    if (pool_allocs || pool_reuse) {
        memory_pool_get_stats(etcp->instance->pkt_pool, pool_allocs, pool_reuse);
    }
}

// ... (rest of functions from master, adapted to ETCP_CONN and using etcp_link_send for tx)

// For tx: assume primary channel (etcp->channels), or iterate if multi-path
static void send_etcp_packet(struct ETCP_CONN* etcp, uint8_t* pkt, uint16_t len) {
    if (!etcp->links) {
        DEBUG_WARN(DEBUG_CATEGORY_CONNECTION, "send_etcp_packet: no links available for node %llu", (unsigned long long)etcp->instance->node_id);
        return; // No links
    }
    
    // Send to primary link (or iterate for multi-path)
    char dg_mem[1600];
    struct ETCP_DGRAM* dg=(struct ETCP_DGRAM*)&dg_mem;
    dg->data_len=len;
    dg->noencrypt_len=0;
    dg->link=etcp->links;// first link TODO: добавить отправку по нескольким линкам
    etcp_encrypt_send(dg);
}

static void update_metrics(struct ETCP_CONN* etcp, uint16_t rtt) {
    etcp->rtt_last = rtt;
    // Update averages, jitter, etc.
    // Placeholder: etcp->rtt_avg_10 = (etcp->rtt_avg_10 * 9 + rtt) / 10;
    // etcp->rtt_avg_100 = (etcp->rtt_avg_100 * 99 + rtt) / 100;
}

uint64_t get_current_time_units() {
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return ((uint64_t)tv.tv_sec * 10000ULL) + (tv.tv_usec / 100);
}

uint16_t get_current_timestamp() {
    return (uint16_t)get_current_time_units();
}


static uint16_t timestamp_diff(uint16_t t1, uint16_t t2) {
    if (t1 >= t2) return t1 - t2;
    return (0xFFFF - t2) + t1 + 1; // Wrap around
}

static int id_compare(uint16_t id1, uint16_t id2) {
    int16_t diff = id1 - id2;
    if (diff == 0) return 0;
    return (diff > 0) ? 1 : -1; // Simple for now, add wrap logic if needed
}

static void retransmit_packet(struct ETCP_CONN* etcp, uint16_t id) {
    sent_packet_t* sent = etcp->sent_list;
    while (sent) {
        if (sent->id == id) {
            send_etcp_packet(etcp, sent->data, sent->data_len);
            etcp->retransmissions_count++;
            break;
        }
        sent = sent->next;
    }
}

static void tx_queue_callback(struct ll_queue* q, struct ll_entry* entry, void* arg) {
    (void)q;
    struct ETCP_CONN* etcp = (struct ETCP_CONN*)arg;
    // Process entry to send packet
    // Placeholder: uint8_t* data = ll_entry_data(entry);
    // size_t len = ll_entry_size(entry);
    // // Build packet and send
    // send_etcp_packet(etcp, data, len);
    queue_resume_callback(q);
}

static void schedule_ack_timer(struct ETCP_CONN* etcp) {
    // Placeholder: set timeout to send ACKs
}

static void request_retransmission_for_gaps(struct ETCP_CONN* etcp) {
    // Detect gaps in rx_list and add to pending_retransmit
}

static void retransmit_check(struct ETCP_CONN* etcp) {
    // Check sent_list for timeouts
}

static void tx_timer_callback(void* arg) {
    struct ETCP_CONN* etcp = (struct ETCP_CONN*)arg;
    // Handle next tx
}

static void retransmit_timer_callback(void* arg) {
    struct ETCP_CONN* etcp = (struct ETCP_CONN*)arg;
    retransmit_check(etcp);
}