#include <stdio.h>
#include <stdint.h>
#include <stdbool.h>
#include <string.h>
#include <time.h>
#include "iperf.h"
#include "bsp_tim.h"
#include "cJSON.h"
#include "socket.h"


static uint8_t cookie[COOKIE_SIZE] = {0};

extern uint8_t ethernet_buf_tx[ETHERNET_BUF_MAX_SIZE];
extern uint8_t ethernet_buf_rx[ETHERNET_BUF_MAX_SIZE];

// Stats结构体初始化函数
void stats_init(Stats *stats, uint32_t pacing_timer_ms)
{
    stats->pacing_timer_ms = pacing_timer_ms;
    stats->running         = false;
    stats->t0              = 0;
    stats->t1              = 0;
    stats->t3              = 0;
    stats->nb0             = 0;
    stats->nb1             = 0;
    stats->np0             = 0;
    stats->np1             = 0;
}

// Stats 开始函数
void stats_start(Stats *stats)
{
    stats->running = true;
    stats->t0 = stats->t1 = get_time_ms();
    stats->nb0 = stats->nb1 = 0;
    stats->np0 = stats->np1 = 0;
    printf("Interval           Transfer     Bitrate\n");
}

// Stats更新函数
void stats_update(Stats *stats, bool final)
{
    if (!stats->running) return;

    uint32_t t2 = get_time_ms();
    uint32_t dt = t2 - stats->t1; //最后更新后的经过时间

    if (final || dt > stats->pacing_timer_ms)
    {
        double ta             = (stats->t1 - stats->t0) / 1e3;       // 之前的时间间隔开始
        double tb             = (t2 - stats->t0) / 1e3;              // 当前时间间隔结束
        double transfer_mbits = (stats->nb1 * 8) / 1e6 / (dt / 1e3); // 兆比特每秒（Mbps）计算

        printf("%5.2f-%-5.2f sec %8u Bytes  %5.2f Mbits/sec\n",
               ta, tb, stats->nb1, transfer_mbits);

        stats->t1  = t2; // 计时器更新为当前时间
        stats->nb1 = 0;  // 周期内数据字节数重置
        stats->np1 = 0;  // 周期内数据包数重置
    }
}

// Stats 停止函数
void stats_stop(Stats *stats)
{
    if (!stats->running) return;

    stats_update(stats, true); // 最后更新
    stats->running = false;

    stats->t3               = get_time_ms(); // 记录结束时间
    uint32_t total_time_ms  = stats->t3 - stats->t0;
    double   total_time_s   = total_time_ms / 1e3;
    double   transfer_mbits = (stats->nb0 * 8) / 1e6 / total_time_s;

    printf("------------------------------------------------------------\n");
    printf("Total: %5.2f sec %8lu Bytes  %5.2f Mbits/sec\n",
           total_time_s, stats->nb0, transfer_mbits);
}

// Stats数据添加函数
void stats_add_bytes(Stats *stats, uint32_t n)
{
    if (!stats->running) return;

    stats->nb0 += n; // 总字节数增加
    stats->nb1 += n; // 周期别字节数增加
    stats->np0 += 1; // 总数据包数量增加
    stats->np1 += 1; // 周期别数据包数量增加
}

void handle_param_exchange(uint8_t socket_ctrl, bool *reverse, bool *udp)
{
    char     buffer[512] = {0};
    uint8_t  cmd;
    uint16_t len        = 0;
    uint8_t  raw_len[4] = {0};
    int      cookie_len;
    cJSON   *json;
    cJSON   *reverseItem;
    cJSON   *udpItem;

    cookie_len = recv(socket_ctrl, cookie, COOKIE_SIZE);
    if (cookie_len != COOKIE_SIZE)
    {
        printf("[iperf] Failed to receive cookie. Received: %d bytes\n", cookie_len);
        return;
    }
    printf("[iperf] Received cookie: %s\n", cookie);

    cmd = PARAM_EXCHANGE;
    send(socket_ctrl, &cmd, 1);
    recv(socket_ctrl, raw_len, 4);

    len = (raw_len[0] << 24) | (raw_len[1] << 16) | (raw_len[2] << 8) | raw_len[3];
    printf("[iperf] Raw length bytes: 0x%02X 0x%02X 0x%02X 0x%02X, Parsed length: %d\n",
           raw_len[0], raw_len[1], raw_len[2], raw_len[3], len);

    recv(socket_ctrl, (uint8_t *)buffer, len);
    buffer[len] = '\0'; // Null-terminate

    printf("[iperf] Received parameters: %s\n", buffer);

    json = cJSON_Parse(buffer);
    if (json == NULL)
    {
        printf("[iperf] Failed to parse JSON: %s\n", cJSON_GetErrorPtr());
    }
    else
    {
        printf("[iperf] Parsed JSON: %s\n", cJSON_Print(json));
        reverseItem = cJSON_GetObjectItem(json, "reverse");
        udpItem     = cJSON_GetObjectItem(json, "udp");

        *reverse = (reverseItem && cJSON_IsBool(reverseItem)) ? reverseItem->valueint : 0;
        *udp     = (udpItem && cJSON_IsBool(udpItem)) ? udpItem->valueint : 0;
        cJSON_Delete(json);
        printf("[iperf] Parsed JSON: reverse=%d, udp=%d\n", *reverse, *udp);
    }
}

void handle_create_streams(uint8_t socket_ctrl, bool udp, uint8_t *remote_ip, uint16_t *remote_port)
{
    uint8_t cmd = CREATE_STREAMS;
    uint8_t received;
	  uint8_t udp_data[5];

    
    if (udp)
    {
        socket(SOCKET_DATA, Sn_MR_UDP, PORT_IPERF, 0x00);

        // Wait for client to connect to data socket
        while (getSn_SR(SOCKET_DATA) != SOCK_UDP)
        {
            if (getSn_SR(SOCKET_DATA) == SOCK_CLOSED)
            {
                printf("[iperf] Data socket closed unexpectedly.\n");
                return;
            }
        }
        send(socket_ctrl, &cmd, 1);
        printf("[iperf] Sent CREATE_STREAMS command.\n");
        // Receive cookie on data socket
        while(getSn_RX_RSR(SOCKET_DATA) == 0)
        {
            // Wait for data to arrive
        }
		received = recvfrom(SOCKET_DATA, udp_data, 4, remote_ip, remote_port);
        udp_data[4] = '\0';
        printf("remote_ip:%d.%d.%d.%d,remote_port:%d,data:%s\r\n",remote_ip[0],remote_ip[1],remote_ip[2],remote_ip[3],*remote_port,udp_data);
        sendto(SOCKET_DATA, "6789", 4, remote_ip, *remote_port);
        printf("[iperf] Sent cookie on UDP data socket.\n");
    }
    else
    {
        send(socket_ctrl, &cmd, 1);
        printf("[iperf] Sent CREATE_STREAMS command.\n");
        socket(SOCKET_DATA, Sn_MR_TCP, PORT_IPERF, Sn_MR_ND);
        listen(SOCKET_DATA);

        // Wait for client to connect to data socket
        while (getSn_SR(SOCKET_DATA) != SOCK_ESTABLISHED)
        {
            if (getSn_SR(SOCKET_DATA) == SOCK_CLOSED)
            {
                printf("[iperf] Data socket closed unexpectedly.\n");
                return;
            }
        }
        printf("[iperf] Data connection established.\n");

        // Receive cookie on data socket
        received = recv(SOCKET_DATA, cookie, COOKIE_SIZE);
        if (received > 0)
        {
            printf("[iperf] Received data cookie: %s\n", cookie);
        }
    }
}

void start_iperf_test(uint8_t socket_ctrl, uint8_t socket_data, Stats *stats, bool reverse, bool udp, uint8_t *remote_ip, uint16_t *remote_port)
{
    bool     running     = true;
    uint8_t  cmd         = 0;

    uint32_t pack_len    = 0;
    uint16_t sent_size, recv_size;

    printf("[iperf] Starting data stream test...\n");

    // Start test
    cmd = TEST_START;
    send(socket_ctrl, &cmd, 1);

    // Running test
    cmd = TEST_RUNNING;
    send(socket_ctrl, &cmd, 1);

    stats_start(stats);

    while (running)
    {
        if (getSn_RX_RSR(SOCKET_CTRL) > 0)
        {
            recv(SOCKET_CTRL, &cmd, 1);
            if (cmd == TEST_END)
            {
                printf("[iperf] TEST_END command received. Stopping test...\n");
                running = false;
                break;
            }
        }
        if (reverse)
        {
            if(udp)
            {
                sent_size = sendto(socket_data, ethernet_buf_tx, ETHERNET_BUF_MAX_SIZE, remote_ip, *(remote_port));
            }
            else
            {
                sent_size = send(socket_data, ethernet_buf_tx, ETHERNET_BUF_MAX_SIZE / 2);
            }
            stats_add_bytes(stats, sent_size);
        }
        else
        {
            getsockopt(socket_data, SO_RECVBUF, &pack_len);
            if (pack_len > 0)
            {
                if(udp)
                {
                    recv_size = recvfrom(socket_data, (uint8_t *)ethernet_buf_rx, ETHERNET_BUF_MAX_SIZE, remote_ip, remote_port); // more fast
                }
                else
                {
                    recv_size = recv(socket_data, (uint8_t *)ethernet_buf_rx, ETHERNET_BUF_MAX_SIZE); // more fast
                }
                stats_add_bytes(stats, recv_size);
            }
            else if (pack_len == 0)
            {
                stats_update(stats, false);
            }
            else
            {
                printf("[iperf] Error during data reception\n");
                break;
            }
        }
        stats_update(stats, false);
    }
    stats_stop(stats);

    exchange_results(SOCKET_CTRL, stats);
}

void exchange_results(uint8_t socket_ctrl, Stats *stats)
{
    uint8_t  cmd        = EXCHANGE_RESULTS;
    uint32_t result_len = 0;
    uint8_t  length_bytes[4];
    char     buffer[1024];
    char    *results_str;
    uint32_t results_len;
    cJSON   *results;
    cJSON   *streams;
    cJSON   *stream;

    // Ask to exchange results
    send(socket_ctrl, &cmd, 1);
    printf("[iperf] Sent EXCHANGE_RESULTS command.\n");

    // Receive client results
    recv(socket_ctrl, (uint8_t *)&result_len, 4);
    result_len = (result_len << 24) | ((result_len << 8) & 0x00FF0000) | ((result_len >> 8) & 0x0000FF00) | (result_len >> 24); // Convert to host-endian

    if (result_len > sizeof(buffer))
    {
        printf("[iperf] Received result length exceeds buffer size.\n");
        return;
    }

    recv(socket_ctrl, (uint8_t *)buffer, result_len);
    buffer[result_len] = '\0'; // Null-terminate the received JSON data
    printf("[iperf] Client results received: %s\n", buffer);

    // Prepare server results
    results = cJSON_CreateObject();
    cJSON_AddNumberToObject(results, "cpu_util_total", 1);
    cJSON_AddNumberToObject(results, "cpu_util_user", 0.5);
    cJSON_AddNumberToObject(results, "cpu_util_system", 0.5);
    cJSON_AddNumberToObject(results, "sender_has_retransmits", 1);
    cJSON_AddStringToObject(results, "congestion_used", "cubic");

    // Streams object
    streams = cJSON_CreateArray();
    stream  = cJSON_CreateObject();
    cJSON_AddNumberToObject(stream, "id", 1);
    cJSON_AddNumberToObject(stream, "bytes", stats->nb0);
    cJSON_AddNumberToObject(stream, "retransmits", 0);
    cJSON_AddNumberToObject(stream, "jitter", 0);
    cJSON_AddNumberToObject(stream, "errors", 0);
    cJSON_AddNumberToObject(stream, "packets", stats->np0);                                // ? ?? ? ??
    cJSON_AddNumberToObject(stream, "start_time", 0);
    cJSON_AddNumberToObject(stream, "end_time", (double)(stats->t3 - stats->t0) / 1000.0); // ?? ?? ??
    cJSON_AddItemToArray(streams, stream);
    cJSON_AddItemToObject(results, "streams", streams);

    // Serialize JSON to string
    results_str = cJSON_PrintUnformatted(results);
    results_len = strlen(results_str);

    // Send server results
    length_bytes[0] = (results_len >> 24) & 0xFF;
    length_bytes[1] = (results_len >> 16) & 0xFF;
    length_bytes[2] = (results_len >> 8) & 0xFF;
    length_bytes[3] = results_len & 0xFF;

    send(socket_ctrl, length_bytes, 4);
    send(socket_ctrl, (uint8_t *)results_str, results_len);

    printf("[iperf] Server results sent.\n");

    cJSON_Delete(results);

    // Ask to display results
    cmd = DISPLAY_RESULTS;
    send(socket_ctrl, &cmd, 1);

    // Wait for IPERF_DONE command
    recv(socket_ctrl, &cmd, 1);
    if (cmd == IPERF_DONE)
    {
        printf("[iperf] Test completed successfully.\n");
    }
    else
    {
        printf("[iperf] Unexpected command received: %d\n", cmd);
    }
}
