#include <stdio.h>
#include <string.h>
#include <inttypes.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/event_groups.h"
#include "esp_log.h"
#include "esp_err.h"
#include "iperf.h"

static const char *TAG = "iperf";

typedef struct {
    iperf_cfg_t cfg;
    bool finish;
    uint32_t total_len;
    uint32_t buffer_len;
    uint8_t *buffer;
    uint32_t sockfd;
} iperf_ctrl_t;

static iperf_ctrl_t s_iperf_ctrl = {0};
static TaskHandle_t s_iperf_task_handle = NULL;

static void socket_send(int sockfd, const uint8_t *buffer, int len)
{
    int actual_send = 0;
    while (actual_send < len) {
        int send_len = send(sockfd, buffer + actual_send, len - actual_send, 0);
        if (send_len < 0) {
            ESP_LOGE(TAG, "send failed: errno %d", errno);
            break;
        }
        actual_send += send_len;
    }
}

static int socket_recv(int sockfd, uint8_t *buffer, int len, TickType_t timeout_ticks)
{
    struct timeval timeout;
    timeout.tv_sec = timeout_ticks / configTICK_RATE_HZ;
    timeout.tv_usec = (timeout_ticks % configTICK_RATE_HZ) * (1000000 / configTICK_RATE_HZ);
    
    if (setsockopt(sockfd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout)) < 0) {
        ESP_LOGE(TAG, "setsockopt failed: errno %d", errno);
        return -1;
    }
    
    return recv(sockfd, buffer, len, 0);
}

static esp_err_t iperf_start_tcp_server(iperf_ctrl_t *ctrl)
{
    struct sockaddr_in addr;
    int listen_sock;
    int opt = 1;

    listen_sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (listen_sock < 0) {
        ESP_LOGE(TAG, "Unable to create socket: errno %d", errno);
        return ESP_FAIL;
    }

    setsockopt(listen_sock, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    addr.sin_family = AF_INET;
    addr.sin_port = htons(ctrl->cfg.sport);
    addr.sin_addr.s_addr = htonl(INADDR_ANY);

    if (bind(listen_sock, (struct sockaddr *)&addr, sizeof(addr)) != 0) {
        ESP_LOGE(TAG, "Socket bind failed: errno %d", errno);
        close(listen_sock);
        return ESP_FAIL;
    }

    if (listen(listen_sock, 5) != 0) {
        ESP_LOGE(TAG, "Socket listen failed: errno %d", errno);
        close(listen_sock);
        return ESP_FAIL;
    }

    ESP_LOGI(TAG, "TCP server listening on port %d", ctrl->cfg.sport);

    while (!ctrl->finish) {
        struct sockaddr_in client_addr;
        socklen_t addr_len = sizeof(client_addr);
        
        int client_sock = accept(listen_sock, (struct sockaddr *)&client_addr, &addr_len);
        if (client_sock < 0) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                continue;
            }
            ESP_LOGE(TAG, "Accept failed: errno %d", errno);
            break;
        }

        ESP_LOGI(TAG, "Client connected from %s", inet_ntoa(client_addr.sin_addr));

        uint64_t total_len = 0;
        int recv_len;

        while (!ctrl->finish) {
            recv_len = socket_recv(client_sock, ctrl->buffer, ctrl->buffer_len, 
                                  pdMS_TO_TICKS(IPERF_SOCKET_RX_TIMEOUT * 1000));
            if (recv_len <= 0) {
                break;
            }
            total_len += recv_len;
        }

        ESP_LOGI(TAG, "TCP server received: %" PRIu64 " bytes", total_len);
        close(client_sock);
    }

    close(listen_sock);
    return ESP_OK;
}

static esp_err_t iperf_start_tcp_client(iperf_ctrl_t *ctrl)
{
    struct sockaddr_in addr;
    int sockfd;

    sockfd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (sockfd < 0) {
        ESP_LOGE(TAG, "Unable to create socket: errno %d", errno);
        return ESP_FAIL;
    }

    addr.sin_family = AF_INET;
    addr.sin_port = htons(ctrl->cfg.dport);
    addr.sin_addr.s_addr = htonl(ctrl->cfg.dip);

    ESP_LOGI(TAG, "Connecting to TCP server...");
    if (connect(sockfd, (struct sockaddr *)&addr, sizeof(addr)) != 0) {
        ESP_LOGE(TAG, "Socket connect failed: errno %d", errno);
        close(sockfd);
        return ESP_FAIL;
    }

    ESP_LOGI(TAG, "Connected to TCP server");

    uint64_t total_len = 0;
    uint32_t start_time = xTaskGetTickCount();
    uint32_t end_time = start_time + pdMS_TO_TICKS(ctrl->cfg.time * 1000);

    while (!ctrl->finish && xTaskGetTickCount() < end_time) {
        socket_send(sockfd, ctrl->buffer, ctrl->buffer_len);
        total_len += ctrl->buffer_len;
    }

    uint32_t actual_time = (xTaskGetTickCount() - start_time) / configTICK_RATE_HZ;
    float bandwidth = (float)total_len * 8 / actual_time / 1000000; // Mbps

    ESP_LOGI(TAG, "TCP client sent: %" PRIu64 " bytes in %" PRIu32 " seconds (%.2f Mbps)", 
             total_len, actual_time, bandwidth);

    close(sockfd);
    return ESP_OK;
}

static esp_err_t iperf_start_udp_server(iperf_ctrl_t *ctrl)
{
    struct sockaddr_in addr;
    int sockfd;
    int opt = 1;

    sockfd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sockfd < 0) {
        ESP_LOGE(TAG, "Unable to create socket: errno %d", errno);
        return ESP_FAIL;
    }

    setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    addr.sin_family = AF_INET;
    addr.sin_port = htons(ctrl->cfg.sport);
    addr.sin_addr.s_addr = htonl(INADDR_ANY);

    if (bind(sockfd, (struct sockaddr *)&addr, sizeof(addr)) != 0) {
        ESP_LOGE(TAG, "Socket bind failed: errno %d", errno);
        close(sockfd);
        return ESP_FAIL;
    }

    ESP_LOGI(TAG, "UDP server listening on port %d", ctrl->cfg.sport);

    uint64_t total_len = 0;
    uint32_t packet_count = 0;
    int32_t last_id = -1;
    uint32_t lost_packets = 0;

    while (!ctrl->finish) {
        int recv_len = socket_recv(sockfd, ctrl->buffer, ctrl->buffer_len,
                                  pdMS_TO_TICKS(IPERF_SOCKET_RX_TIMEOUT * 1000));
        if (recv_len <= 0) {
            continue;
        }

        total_len += recv_len;
        packet_count++;

        // Check for lost packets using UDP header
        if (recv_len >= sizeof(udp_datagram)) {
            udp_datagram *header = (udp_datagram *)ctrl->buffer;
            int32_t current_id = ntohl(header->id);
            
            if (last_id >= 0 && current_id > last_id + 1) {
                lost_packets += (current_id - last_id - 1);
            }
            last_id = current_id;
        }
    }

    float loss_rate = packet_count > 0 ? (float)lost_packets * 100 / packet_count : 0;
    ESP_LOGI(TAG, "UDP server received: %" PRIu64 " bytes, %" PRIu32 " packets, %" PRIu32 " lost (%.2f%%)", 
             total_len, packet_count, lost_packets, loss_rate);

    close(sockfd);
    return ESP_OK;
}

static esp_err_t iperf_start_udp_client(iperf_ctrl_t *ctrl)
{
    struct sockaddr_in addr;
    int sockfd;

    sockfd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sockfd < 0) {
        ESP_LOGE(TAG, "Unable to create socket: errno %d", errno);
        return ESP_FAIL;
    }

    addr.sin_family = AF_INET;
    addr.sin_port = htons(ctrl->cfg.dport);
    addr.sin_addr.s_addr = htonl(ctrl->cfg.dip);

    ESP_LOGI(TAG, "Starting UDP client");

    uint64_t total_len = 0;
    uint32_t packet_count = 0;
    uint32_t start_time = xTaskGetTickCount();
    uint32_t end_time = start_time + pdMS_TO_TICKS(ctrl->cfg.time * 1000);

    while (!ctrl->finish && xTaskGetTickCount() < end_time) {
        udp_datagram *header = (udp_datagram *)ctrl->buffer;
        header->id = htonl(packet_count);
        
        int send_len = sendto(sockfd, ctrl->buffer, ctrl->buffer_len, 0,
                             (struct sockaddr *)&addr, sizeof(addr));
        if (send_len > 0) {
            total_len += send_len;
            packet_count++;
        }

        // Bandwidth limiting
        if (ctrl->cfg.bw_lim > 0) {
            vTaskDelay(1);
        }
    }

    uint32_t actual_time = (xTaskGetTickCount() - start_time) / configTICK_RATE_HZ;
    float bandwidth = actual_time > 0 ? (float)total_len * 8 / actual_time / 1000000 : 0;

    ESP_LOGI(TAG, "UDP client sent: %" PRIu64 " bytes, %" PRIu32 " packets in %" PRIu32 " seconds (%.2f Mbps)", 
             total_len, packet_count, actual_time, bandwidth);

    close(sockfd);
    return ESP_OK;
}

static void iperf_task(void *arg)
{
    iperf_ctrl_t *ctrl = (iperf_ctrl_t *)arg;

    if (ctrl->cfg.flag & IPERF_FLAG_TCP) {
        if (ctrl->cfg.flag & IPERF_FLAG_SERVER) {
            iperf_start_tcp_server(ctrl);
        } else {
            iperf_start_tcp_client(ctrl);
        }
    } else {
        if (ctrl->cfg.flag & IPERF_FLAG_SERVER) {
            iperf_start_udp_server(ctrl);
        } else {
            iperf_start_udp_client(ctrl);
        }
    }

    if (ctrl->buffer) {
        free(ctrl->buffer);
        ctrl->buffer = NULL;
    }

    ESP_LOGI(TAG, "iperf task finished");
    s_iperf_task_handle = NULL;
    vTaskDelete(NULL);
}

void iperf_start(iperf_cfg_t *cfg)
{
    if (s_iperf_task_handle != NULL) {
        ESP_LOGW(TAG, "iperf is already running");
        return;
    }

    memcpy(&s_iperf_ctrl.cfg, cfg, sizeof(iperf_cfg_t));
    s_iperf_ctrl.finish = false;
    
    // Allocate buffer
    if (cfg->flag & IPERF_FLAG_TCP) {
        s_iperf_ctrl.buffer_len = cfg->flag & IPERF_FLAG_SERVER ? 
                                  IPERF_TCP_RX_LEN : IPERF_TCP_TX_LEN;
    } else {
        s_iperf_ctrl.buffer_len = cfg->flag & IPERF_FLAG_SERVER ? 
                                  IPERF_UDP_RX_LEN : IPERF_UDP_TX_LEN;
    }

    s_iperf_ctrl.buffer = (uint8_t *)malloc(s_iperf_ctrl.buffer_len);
    if (!s_iperf_ctrl.buffer) {
        ESP_LOGE(TAG, "Failed to allocate buffer");
        return;
    }

    memset(s_iperf_ctrl.buffer, 0, s_iperf_ctrl.buffer_len);

    xTaskCreate(iperf_task, "iperf", 4096, &s_iperf_ctrl, IPERF_TRAFFIC_TASK_PRIORITY, 
                &s_iperf_task_handle);
}

void iperf_stop(void)
{
    if (s_iperf_task_handle != NULL) {
        s_iperf_ctrl.finish = true;
        ESP_LOGI(TAG, "Stopping iperf...");
    }
}