#include "etcp_connections.h"
#include <arpa/inet.h>
#include "debug_config.h"
#include "routing.h"
#include "utun_instance.h"
#include "config_parser.h"
#include "crc32.h"
#include "etcp.h"
#include <stdlib.h>
#include <string.h>
#include <time.h>

// Бинарный поиск линка по ip_port_hash
static int find_link_index(struct ETCP_CONNECTIONS* conns, uint32_t hash) {
    if (!conns || conns->num_channels == 0) return -1;
    
    int left = 0;
    int right = conns->num_channels - 1;
    
    while (left <= right) {
        int mid = left + (right - left) / 2;
        if (conns->links[mid]->ip_port_hash == hash) {
            return mid;
        } else if (conns->links[mid]->ip_port_hash < hash) {
            left = mid + 1;
        } else {
            right = mid - 1;
        }
    }
    
    return -(left + 1);
}

// Реалокация массива линков с увеличением в 2 раза
static int realloc_links(struct ETCP_CONNECTIONS* conns) {
    size_t new_max = conns->max_channels == 0 ? 8 : conns->max_channels * 2;
    struct ETCP_LINK** new_links = realloc(conns->links, new_max * sizeof(struct ETCP_LINK*));
    if (!new_links) return -1;
    
    conns->links = new_links;
    conns->max_channels = new_max;
    return 0;
}

// Вставка линка в отсортированный массив
static int insert_link(struct ETCP_CONNECTIONS* conns, struct ETCP_LINK* link) {
    if (!conns || !link) return -1;
    
    if (conns->num_channels >= conns->max_channels) {
        if (realloc_links(conns) < 0) return -1;
    }
    
    int idx = find_link_index(conns, link->ip_port_hash);
    if (idx >= 0) return -1;
    
    idx = -(idx + 1);
    
    if (idx < (int)conns->num_channels) {
        memmove(&conns->links[idx + 1], &conns->links[idx], 
                (conns->num_channels - idx) * sizeof(struct ETCP_LINK*));
    }
    
    conns->links[idx] = link;
    conns->num_channels++;
    return 0;
}

// Удаление линка из массива
static void remove_link(struct ETCP_CONNECTIONS* conns, uint32_t hash) {
    if (!conns || conns->num_channels == 0) return;
    
    int idx = find_link_index(conns, hash);
    if (idx < 0) return;
    
    if (idx < (int)conns->num_channels - 1) {
        memmove(&conns->links[idx], &conns->links[idx + 1], 
                (conns->num_channels - idx - 1) * sizeof(struct ETCP_LINK*));
    }
    
    conns->num_channels--;
}

struct ETCP_CONNECTIONS* etcp_connections_init(struct ETCP_SOCKET* socket) {
    if (!socket) return NULL;
    
    struct ETCP_CONNECTIONS* conns = calloc(1, sizeof(struct ETCP_CONNECTIONS));
    if (!conns) return NULL;
    
    memcpy(&conns->socket, socket, sizeof(struct ETCP_SOCKET));
    return conns;
}

void etcp_connections_destroy(struct ETCP_CONNECTIONS* conns) {
    if (!conns) return;
    
    for (size_t i = 0; i < conns->num_channels; i++) {
        etcp_link_close(conns->links[i]);
    }
    
    free(conns->links);
    free(conns);
}

struct ETCP_LINK* etcp_link_new(struct ETCP_CONN* etcp, struct ETCP_SOCKET* socket, 
                               const struct sockaddr* remote_addr, socklen_t addr_len) {
    if (!socket || !remote_addr || addr_len == 0) return NULL;
    
    struct ETCP_LINK* link = calloc(1, sizeof(struct ETCP_LINK));
    if (!link) return NULL;
    
    link->socket = (struct ETCP_SOCKET*)socket;
    link->etcp = etcp;
    memcpy(&link->remote_addr, remote_addr, addr_len);
    link->remote_addr_len = addr_len;
    link->last_activity = time(NULL);
    
    uint8_t* addr_bytes = (uint8_t*)remote_addr;
    link->ip_port_hash = crc32_calc(addr_bytes, addr_len);
    
    return link;
}

void etcp_link_close(struct ETCP_LINK* link) {
    if (!link) return;
    free(link);
}

int etcp_input(struct packet_buffer* pkt, struct ETCP_SOCKET* socket, struct ETCP_CONNECTIONS* conns) {
    if (!pkt || !socket || !conns) return -1;
    
    if (pkt->metadata.data_len < 1) return -1;
    
    uint8_t* data = pkt->data;
    uint8_t cmd = data[0];
    
    struct sockaddr* src_addr = packet_remote_addr(pkt);
    socklen_t addr_len = sizeof(struct sockaddr_storage);
    
    if (cmd == ETCP_INIT_REQUEST) {
        if (pkt->metadata.data_len < 1 + 8 + SC_PUBKEY_SIZE + 2 + 2) {
            return -1;
        }
        
        struct ETCP_LINK* link = etcp_link_find_by_addr(conns, src_addr, addr_len);
        if (link) {
            if (link->initialized) {
                etcp_conn_reset(link->etcp);
            }
            etcp_link_remove_from_connections(conns, link);
            etcp_link_close(link);
        }
        
        link = etcp_link_new(NULL, socket, src_addr, addr_len);
        if (!link) return -1;
        
        uint8_t* p = &data[1];
        link->peer_node_id = 0;
        memcpy(&link->peer_node_id, p, 8);
        p += 8;
        
        memcpy(link->peer_public_key, p, SC_PUBKEY_SIZE);
        link->has_peer_key = 1;
        p += SC_PUBKEY_SIZE;
        
        uint16_t peer_mtu = ntohs(*(uint16_t*)p);
        p += 2;
        uint16_t peer_keepalive = ntohs(*(uint16_t*)p);
        
        link->mtu = peer_mtu;
        link->keepalive_interval = peer_keepalive;
        
        etcp_link_add_to_connections(conns, link);
        
        return etcp_link_send_init_response(link, 1500, 30);
    }
    
    struct ETCP_LINK* link = etcp_link_find_by_addr(conns, src_addr, addr_len);
    if (!link) {
        struct ETCP_LINK* temp_link = etcp_link_new(NULL, socket, src_addr, addr_len);
        if (temp_link) {
            etcp_link_send_reset(temp_link);
            etcp_link_close(temp_link);
        }
        return -1;
    }
    
    if (!link->initialized) {
        switch (cmd) {
            case ETCP_INIT_RESPONSE:
                if (pkt->metadata.data_len < 1 + 8 + 2 + 2) return -1;
                
                link->peer_node_id = 0;
                memcpy(&link->peer_node_id, &data[1], 8);
                
                link->mtu = ntohs(*(uint16_t*)&data[1 + 8]);
                link->keepalive_interval = ntohs(*(uint16_t*)&data[1 + 8 + 2]);
                link->initialized = 1;
                break;
                
            case ETCP_CHANNEL_INIT:
                if (pkt->metadata.data_len < 1 + 8 + 2) return -1;
                
                link->peer_node_id = 0;
                memcpy(&link->peer_node_id, &data[1], 8);
                link->keepalive_interval = ntohs(*(uint16_t*)&data[1 + 8]);
                
                return etcp_link_send_channel_response(link);
                
            case ETCP_CHANNEL_RESPONSE:
                if (pkt->metadata.data_len < 1 + 8) return -1;
                
                link->peer_node_id = 0;
                memcpy(&link->peer_node_id, &data[1], 8);
                link->initialized = 1;
                break;
                
            case ETCP_RESET:
                etcp_link_remove_from_connections(conns, link);
                etcp_link_close(link);
                break;
                
            default:
                etcp_link_send_reset(link);
                return -1;
        }
    } else {
        if (cmd >= 0x02 && cmd <= 0x06) {
            etcp_link_send_reset(link);
            return -1;
        }
    }
    
    link->last_activity = time(NULL);
    
    return 0;
}

int etcp_link_init(struct ETCP_CONN* etcp, struct ETCP_LINK* link) { (void)etcp; (void)link;
    return 0;
}

int etcp_link_send(struct ETCP_CONN* etcp, struct ETCP_LINK* link, 
                  const uint8_t* data, size_t len) {
    if (!etcp || !link || !data || len == 0) return -1;
    
    link->last_activity = time(NULL);
    
    return etcp_socket_send((struct ETCP_SOCKET*)link->socket, data, len, 
                           (struct sockaddr*)&link->remote_addr);
}

int etcp_link_send_init(struct ETCP_LINK* link, int mtu, uint16_t keepalive_interval) {
    if (!link || mtu <= 0 || mtu > 65535) return -1;
    
    uint8_t packet[1 + 8 + SC_PUBKEY_SIZE + 2 + 2];
    uint8_t* p = packet;
    
    *p++ = ETCP_INIT_REQUEST;
    
    uint64_t node_id = link->etcp->node_id;
    memcpy(p, &node_id, 8);
    p += 8;
    
    sc_context_t* sc = link->etcp->crypto_ctx;
    if (!sc || !sc->initialized) return -1;
    
    memcpy(p, sc->public_key, SC_PUBKEY_SIZE);
    p += SC_PUBKEY_SIZE;
    
    uint16_t mtu_be = htons(mtu);
    memcpy(p, &mtu_be, 2);
    p += 2;
    
    uint16_t keepalive_be = htons(keepalive_interval);
    memcpy(p, &keepalive_be, 2);
    p += 2;
    
    link->mtu = mtu;
    link->keepalive_interval = keepalive_interval;
    
    return etcp_socket_send((struct ETCP_SOCKET*)link->socket, packet, sizeof(packet), 
                           (struct sockaddr*)&link->remote_addr);
}

int etcp_link_send_init_response(struct ETCP_LINK* link, int mtu, uint16_t keepalive_interval) {
    if (!link || mtu <= 0 || mtu > 65535) return -1;
    
    uint8_t packet[1 + 8 + 2 + 2];
    uint8_t* p = packet;
    
    *p++ = ETCP_INIT_RESPONSE;
    
    uint64_t node_id = link->etcp->node_id;
    memcpy(p, &node_id, 8);
    p += 8;
    
    uint16_t mtu_be = htons(mtu);
    memcpy(p, &mtu_be, 2);
    p += 2;
    
    uint16_t keepalive_be = htons(keepalive_interval);
    memcpy(p, &keepalive_be, 2);
    p += 2;
    
    link->mtu = mtu;
    link->keepalive_interval = keepalive_interval;
    link->initialized = 1;
    
    return etcp_socket_send((struct ETCP_SOCKET*)link->socket, packet, sizeof(packet), 
                           (struct sockaddr*)&link->remote_addr);
}

int etcp_link_send_channel_init(struct ETCP_LINK* link, uint16_t keepalive_interval) {
    if (!link) return -1;
    
    uint8_t packet[1 + 8 + 2];
    uint8_t* p = packet;
    
    *p++ = ETCP_CHANNEL_INIT;
    
    uint64_t node_id = link->etcp->node_id;
    memcpy(p, &node_id, 8);
    p += 8;
    
    uint16_t keepalive_be = htons(keepalive_interval);
    memcpy(p, &keepalive_be, 2);
    p += 2;
    
    link->keepalive_interval = keepalive_interval;
    
    return etcp_socket_send((struct ETCP_SOCKET*)link->socket, packet, sizeof(packet), 
                           (struct sockaddr*)&link->remote_addr);
}

int etcp_link_send_channel_response(struct ETCP_LINK* link) {
    if (!link) return -1;
    
    uint8_t packet[1 + 8];
    uint8_t* p = packet;
    
    *p++ = ETCP_CHANNEL_RESPONSE;
    
    uint64_t node_id = link->etcp->node_id;
    memcpy(p, &node_id, 8);
    p += 8;
    
    link->initialized = 1;
    
    return etcp_socket_send((struct ETCP_SOCKET*)link->socket, packet, sizeof(packet), 
                           (struct sockaddr*)&link->remote_addr);
}

int etcp_link_send_reset(struct ETCP_LINK* link) {
    if (!link) return -1;
    
    uint8_t packet[1] = {ETCP_RESET};
    
    return etcp_socket_send((struct ETCP_SOCKET*)link->socket, packet, sizeof(packet), 
                           (struct sockaddr*)&link->remote_addr);
}

struct ETCP_LINK* etcp_link_find_by_addr(struct ETCP_CONNECTIONS* conns, 
                                        const struct sockaddr* addr, socklen_t addr_len) {
    if (!conns || !addr || addr_len == 0) return NULL;
    
    uint32_t hash = crc32_calc((const uint8_t*)addr, addr_len);
    int idx = find_link_index(conns, hash);
    
    if (idx >= 0) {
        return conns->links[idx];
    }
    return NULL;
}

int etcp_link_add_to_connections(struct ETCP_CONNECTIONS* conns, struct ETCP_LINK* link) {
    if (!conns || !link) return -1;
    return insert_link(conns, link);
}

void etcp_link_remove_from_connections(struct ETCP_CONNECTIONS* conns, struct ETCP_LINK* link) {
    if (!conns || !link) return;
    remove_link(conns, link->ip_port_hash);
}

// SOCKET FUNCTIONS (moved from etcp_sockets)
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <arpa/inet.h>

struct ETCP_SOCKET* etcp_socket_create(struct ETCP_CONN* etcp, const char* bind_addr, uint16_t port) {
    if (!etcp) return NULL;
    
    struct ETCP_SOCKET* sock = calloc(1, sizeof(struct ETCP_SOCKET));
    if (!sock) return NULL;
    
    sock->etcp = etcp;
    sock->socket_id = (uint32_t)(uintptr_t)sock;
    
    sock->fd = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock->fd < 0) {
        free(sock);
        return NULL;
    }
    
    int flags = fcntl(sock->fd, F_GETFL, 0);
    fcntl(sock->fd, F_SETFL, flags | O_NONBLOCK);
    
    struct sockaddr_in* addr = (struct sockaddr_in*)&sock->local_addr;
    addr->sin_family = AF_INET;
    addr->sin_port = htons(port);
    sock->local_addr_len = sizeof(struct sockaddr_in);
    
    if (bind_addr) {
        inet_pton(AF_INET, bind_addr, &addr->sin_addr);
    } else {
        addr->sin_addr.s_addr = INADDR_ANY;
    }
    
    if (bind(sock->fd, (struct sockaddr*)addr, sock->local_addr_len) < 0) {
        close(sock->fd);
        free(sock);
        return NULL;
    }
    
    if (getsockname(sock->fd, (struct sockaddr*)addr, &sock->local_addr_len) < 0) {
        close(sock->fd);
        free(sock);
        return NULL;
    }
    
    return sock;
}

void etcp_socket_destroy(struct ETCP_SOCKET* sock) {
    if (!sock) return;
    
    if (sock->fd >= 0) {
        close(sock->fd);
    }
    
    free(sock);
}

int etcp_socket_send(struct ETCP_SOCKET* sock, const uint8_t* data, size_t len, const struct sockaddr* to_addr) {
    if (!sock || sock->fd < 0 || !data || !to_addr) {
        return -1;
    }
    
    socklen_t addr_len = (to_addr->sa_family == AF_INET) ? 
                         sizeof(struct sockaddr_in) : 
                         sizeof(struct sockaddr_in6);
    
    return sendto(sock->fd, data, len, 0, to_addr, addr_len);
}

int etcp_socket_get_fd(const struct ETCP_SOCKET* sock) {
    if (!sock) return -1;
    return sock->fd;
}

void etcp_socket_get_local_addr(const struct ETCP_SOCKET* sock, struct sockaddr_storage* addr, socklen_t* addr_len) {
    if (!sock || !addr || !addr_len) return;
    
    memcpy(addr, &sock->local_addr, sock->local_addr_len);
    *addr_len = sock->local_addr_len;
}

int etcp_socket_set_option(struct ETCP_SOCKET* sock, int level, int optname, const void* optval, socklen_t optlen) {
    if (!sock || sock->fd < 0) return -1;
    return setsockopt(sock->fd, level, optname, optval, optlen);
}

// Helper function to parse IP:port address
static int parse_ip_port(const char* addr_str, struct sockaddr_storage* addr, socklen_t* addr_len) {
    if (!addr_str || !addr || !addr_len) return -1;
    
    char* colon = strchr(addr_str, ':');
    if (!colon) return -1;
    
    char ip_str[INET6_ADDRSTRLEN];
    size_t ip_len = colon - addr_str;
    if (ip_len >= sizeof(ip_str)) return -1;
    
    strncpy(ip_str, addr_str, ip_len);
    ip_str[ip_len] = '\0';
    
    char* port_str = colon + 1;
    int port = atoi(port_str);
    if (port <= 0 || port > 65535) return -1;
    
    struct sockaddr_in* sin = (struct sockaddr_in*)addr;
    memset(sin, 0, sizeof(struct sockaddr_in));
    sin->sin_family = AF_INET;
    sin->sin_port = htons(port);
    
    if (inet_pton(AF_INET, ip_str, &sin->sin_addr) != 1) {
        return -1;
    }
    
    *addr_len = sizeof(struct sockaddr_in);
    return 0;
}

// Initialize connections from configuration
int init_connections(struct UTUN_INSTANCE* instance) {
    if (!instance || !instance->config) {
        DEBUG_ERROR(DEBUG_CATEGORY_CONFIG, "Invalid instance or config");
        return -1;
    }
    
    struct utun_config* config = instance->config;
    int total_servers = config->server_count;
    
    if (total_servers == 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_CONFIG, "No servers found in configuration");
        return -1;
    }
    
    instance->connections = calloc(total_servers, sizeof(conn_handle_t*));
    if (!instance->connections) {
        DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "Failed to allocate connections array");
        return -1;
    }
    
    instance->connection_count = 0;
    
    for (int i = 0; i < total_servers; i++) {
        struct server_config* server = &config->servers[i];
        
        // Parse server bind address
        struct sockaddr_storage bind_addr;
        socklen_t bind_addr_len;
        if (parse_ip_port(server->addr, &bind_addr, &bind_addr_len) < 0) {
            DEBUG_ERROR(DEBUG_CATEGORY_CONFIG, "Invalid server address '%s'", server->addr);
            continue;
        }
        
        // Get bind IP and port
        struct sockaddr_in* sin = (struct sockaddr_in*)&bind_addr;
        char bind_ip_str[INET_ADDRSTRLEN];
        if (!inet_ntop(AF_INET, &sin->sin_addr, bind_ip_str, sizeof(bind_ip_str))) {
            DEBUG_ERROR(DEBUG_CATEGORY_CONFIG, "Failed to convert bind address");
            continue;
        }
        uint16_t bind_port = ntohs(sin->sin_port);
        
        // Create ETCP instance
        struct ETCP_CONN* etcp = etcp_create(instance);
        if (!etcp) {
            DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Failed to create ETCP for server %s", server->name);
            continue;
        }
        
        // Create socket
        struct ETCP_SOCKET* sock = etcp_socket_create(etcp, bind_ip_str, bind_port);
        if (!sock) {
            DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Failed to create socket for %s", server->name);
            etcp_destroy(etcp);
            continue;
        }
        
        // Create connections manager
        struct ETCP_CONNECTIONS* conns = etcp_connections_init(sock);
        if (!conns) {
            DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Failed to create connections for %s", server->name);
            etcp_socket_destroy(sock);
            etcp_destroy(etcp);
            continue;
        }
        
        // Create connection handle
        conn_handle_t* handle = calloc(1, sizeof(conn_handle_t));
        if (!handle) {
            DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "Failed to allocate connection handle");
            etcp_connections_destroy(conns);
            etcp_socket_destroy(sock);
            etcp_destroy(etcp);
            continue;
        }
        
        handle->etcp = etcp;
        handle->conns = conns;
        
        // Set first listen socket
        if (i == 0) {
            instance->first_listen_socket = (struct ETCP_SOCKET*)sock;
        }
        
        // Process clients that connect to this server
        for (int j = 0; j < config->client_count; j++) {
            struct client_config* client = &config->clients[j];
            
            // Check if this client is for current server
            if (strcmp(client->from, server->name) != 0) {
                continue;
            }
            
            // Parse client address
            struct sockaddr_storage client_addr;
            socklen_t client_addr_len;
            if (parse_ip_port(client->to_addr, &client_addr, &client_addr_len) < 0) {
                DEBUG_ERROR(DEBUG_CATEGORY_CONFIG, "Invalid client address '%s'", client->to_addr);
                continue;
            }
            
            // Create link for this client
            struct ETCP_LINK* link = etcp_link_new(etcp, sock, 
                                                   (struct sockaddr*)&client_addr, 
                                                   client_addr_len);
            if (!link) {
                DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Failed to create link for client %s", client->name);
                continue;
            }
            
            // Add link to connections
            if (etcp_link_add_to_connections(conns, link) < 0) {
                DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Failed to add link for client %s", client->name);
                etcp_link_close(link);
                continue;
            }
            
            printf("  Added client %s -> %s\n", client->name, client->to_addr);
        }
        
        instance->connections[instance->connection_count++] = handle;
        printf("Initialized server %s on %s (links: %zu)\n", 
               server->name, server->addr, conns->num_channels);
    }
    
    if (instance->connection_count == 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "No connections initialized");
        free(instance->connections);
        instance->connections = NULL;
        return -1;
    }
    
    printf("Initialized %d connections\n", instance->connection_count);
    return 0;
}

void conn_destroy(conn_handle_t* handle) {
    if (!handle) return;
    
    if (handle->conns) {
        etcp_socket_destroy((struct ETCP_SOCKET*)&handle->conns->socket);
        etcp_connections_destroy(handle->conns);
    }
    
    if (handle->etcp) {
        etcp_destroy(handle->etcp);
    }
    
    free(handle);
}

int conn_send(conn_handle_t* conn, const uint8_t* data, size_t len) {
    if (!conn || !conn->etcp || !data || len == 0) {
        return -1;
    }

    struct ETCP_CONN* etcp = conn->etcp;
    struct packet_buffer* pkt = packet_pool_get(&etcp->packet_pool);
    if (!pkt) {
        return -1;
    }

    if (len > 1500) {
        len = 1500;
    }

    memcpy(packet_data(pkt), data, len);
    pkt->metadata.data_len = len;
    
    etcp_conn_input(etcp, pkt);
    return 0;
}



size_t etcp_connections_get_channel_count(const struct ETCP_CONNECTIONS* conns) {
    if (!conns) return 0;
    return conns->num_channels;
}