#include "WIFI.h"

#if 0
static void dump_bytes(const uint8_t *bptr, uint32_t len) {
    unsigned int i = 0;

    printf("dump_bytes %d", len);
    for (i = 0; i < len;) {
        if ((i & 0x0f) == 0) {
            printf("\n");
        } else if ((i & 0x07) == 0) {
            printf(" ");
        }
        printf("%02x ", bptr[i++]);
    }
    printf("\n");
}
#define DUMP_BYTES dump_bytes  // 调试用的字节转储函数
#else
#define DUMP_BYTES(A,B)
#endif

typedef struct TCP_CLIENT_T_ {
    struct tcp_pcb *tcp_pcb;  // TCP控制块指针
    ip_addr_t remote_addr;  // 远程地址
    uint8_t buffer[BUF_SIZE];  // 接收缓冲区
    int buffer_len;  // 缓冲区当前长度
    int sent_len;  // 已发送的数据长度
    bool complete;  // 是否完成测试
    int run_count;  // 运行次数
    bool connected;  // 是否已连接
} TCP_CLIENT_T;

static err_t tcp_client_close(void *arg) {
    TCP_CLIENT_T *state = (TCP_CLIENT_T*)arg;
    err_t err = ERR_OK;
    if (state->tcp_pcb != NULL) {
        tcp_arg(state->tcp_pcb, NULL);  // 清除回调参数
        tcp_poll(state->tcp_pcb, NULL, 0);  // 清除轮询回调
        tcp_sent(state->tcp_pcb, NULL);  // 清除发送回调
        tcp_recv(state->tcp_pcb, NULL);  // 清除接收回调
        tcp_err(state->tcp_pcb, NULL);  // 清除错误回调
        err = tcp_close(state->tcp_pcb);  // 关闭TCP连接
        if (err != ERR_OK) {
            DEBUG_printf("close failed %d, calling abort\n", err);  // 如果关闭失败，调用abort
            tcp_abort(state->tcp_pcb);
            err = ERR_ABRT;
        }
        state->tcp_pcb = NULL;  // 清除TCP控制块指针
    }
    return err;
}

// 操作结果回调
static err_t tcp_result(void *arg, int status) {
    TCP_CLIENT_T *state = (TCP_CLIENT_T*)arg;
    if (status == 0) {
        DEBUG_printf("test success\n");  // 测试成功
    } else {
        DEBUG_printf("test failed %d\n", status);  // 测试失败
    }
    state->complete = true;  // 标记测试完成
    return tcp_client_close(arg);  // 关闭连接
}

// 轮询回调
static err_t tcp_client_poll(void *arg, struct tcp_pcb *tpcb) {
    DEBUG_printf("tcp_client_poll\n");
    return tcp_result(arg, -1);  // 无响应视为错误
}

// 错误处理回调
static void tcp_client_err(void *arg, err_t err) {
    if (err != ERR_ABRT) {
        DEBUG_printf("tcp_client_err %d\n", err);  // 打印错误信息
        tcp_result(arg, err);
    }
}

// 连接建立回调
static err_t tcp_client_connected(void *arg, struct tcp_pcb *tpcb, err_t err) {
    TCP_CLIENT_T *state = (TCP_CLIENT_T*)arg;
    if (err != ERR_OK) {
        printf("connect failed %d\n", err);  // 连接失败
        return tcp_result(arg, err);
    }
    state->connected = true;  // 标记已连接
    DEBUG_printf("连接成功，准备发送消息！！！\n");  // 连接成功，准备发送消息

    // 定义要发送的字符串
    const char *message = "PICO W\n";
    state->buffer_len = strlen(message);  // 设置要发送的字符串长度
    memcpy(state->buffer, message, state->buffer_len);  // 将字符串复制到缓冲区

    // 发送字符串到服务器
    err = tcp_write(tpcb, state->buffer, state->buffer_len, TCP_WRITE_FLAG_COPY);
    if (err != ERR_OK) {
        DEBUG_printf("发送失败: %d\n", err);  // 发送失败
        return tcp_result(arg, -1);
    }
    DEBUG_printf("发送数据: %s\n", message);  // 打印发送的消息

    memset(state->buffer, 0, state->buffer_len);
    state->buffer_len = 0;

    return ERR_OK;
}

err_t tcp_client_recv(void *arg, struct tcp_pcb *tpcb, struct pbuf *p, err_t err) {
    TCP_CLIENT_T *state = (TCP_CLIENT_T*)arg;
    if (!p) {
        DEBUG_printf("接收数据失败，p 为 NULL\n");  // 添加调试信息
        return tcp_result(arg, -1);
    }
    // this method is callback from lwIP, so cyw43_arch_lwip_begin is not required, however you
    // can use this method to cause an assertion in debug mode, if this method is called when
    // cyw43_arch_lwip_begin IS needed
    cyw43_arch_lwip_check();
    if (p->tot_len > 0) {
        DEBUG_printf("recv %d err %d\n", p->tot_len, err);
        for (struct pbuf *q = p; q != NULL; q = q->next) {
            DUMP_BYTES(q->payload, q->len);
        }
        // Receive the buffer
        const uint16_t buffer_left = BUF_SIZE - state->buffer_len;
        state->buffer_len += pbuf_copy_partial(p, state->buffer + state->buffer_len,
                                               p->tot_len > buffer_left ? buffer_left : p->tot_len, 0);
        tcp_recved(tpcb, p->tot_len);
    }

    // 打印当前缓冲区内容
    printf("接收数据: %s", state->buffer);

    pbuf_free(p);

    memset(state->buffer, 0, state->buffer_len);
    state->buffer_len = 0;

    return ERR_OK;
}

// 发送数据回调
static err_t tcp_client_sent(void *arg, struct tcp_pcb *tpcb, u16_t len) {
    TCP_CLIENT_T *state = (TCP_CLIENT_T*)arg;
    DEBUG_printf("tcp_client_sent %u\n", len);  // 打印发送的数据长度
    state->sent_len += len;  // 更新已发送的数据长度

    if (state->sent_len >= BUF_SIZE) {
        state->run_count++;  // 增加运行次数
        if (state->run_count >= TEST_ITERATIONS) {
            tcp_result(arg, 0);  // 达到迭代次数，测试成功
            return ERR_OK;
        }

        // 准备接收新的数据缓冲区
        state->buffer_len = 0;
        state->sent_len = 0;
        DEBUG_printf("Waiting for buffer from server\n");
    }

    return ERR_OK;
}

// 打开TCP连接
static bool tcp_client_open(void *arg) {
    TCP_CLIENT_T *state = (TCP_CLIENT_T*)arg;
    DEBUG_printf("Connecting to %s port %u\n", ip4addr_ntoa(&state->remote_addr), TCP_PORT);
    state->tcp_pcb = tcp_new_ip_type(IP_GET_TYPE(&state->remote_addr));  // 创建TCP控制块
    if (!state->tcp_pcb) {
        DEBUG_printf("failed to create pcb\n");  // 创建失败
        return false;
    }

    tcp_arg(state->tcp_pcb, state);  // 设置回调参数
    tcp_poll(state->tcp_pcb, tcp_client_poll, POLL_TIME_S * 2);  // 设置轮询回调
    tcp_sent(state->tcp_pcb, tcp_client_sent);  // 设置发送回调
    tcp_recv(state->tcp_pcb, tcp_client_recv);  // 设置接收回调
    tcp_err(state->tcp_pcb, tcp_client_err);  // 设置错误回调

    state->buffer_len = 0;  // 初始化缓冲区长度

    // 使用cyw43_arch_lwip_begin/end确保正确锁定
    cyw43_arch_lwip_begin();
    err_t err = tcp_connect(state->tcp_pcb, &state->remote_addr, TCP_PORT, tcp_client_connected);  // 尝试连接
    cyw43_arch_lwip_end();

    return err == ERR_OK;
}

// 初始化TCP客户端
static TCP_CLIENT_T* tcp_client_init(void) {
    TCP_CLIENT_T *state = calloc(1, sizeof(TCP_CLIENT_T));  // 分配内存
    if (!state) {
        DEBUG_printf("failed to allocate state\n");  // 分配失败
        return NULL;
    }

    
    // 如果 TEST_TCP_SERVER_IP 是 IP 地址，直接解析为 IP 地址
    if (ip4addr_aton(TEST_TCP_SERVER_IP, &state->remote_addr) == 0) {
        // 如果解析失败，尝试作为域名解析
        DEBUG_printf("Failed to parse IP address, trying DNS resolution\n");

        err_t err = dns_gethostbyname(TEST_TCP_SERVER_IP, &state->remote_addr, NULL, NULL);
        if (err != ERR_OK) {
            DEBUG_printf("DNS resolution failed\n");
            free(state);
            return NULL;
        }
    }

    return state;
}

// 运行TCP客户端测试
void run_tcp_client_test(void) {
    TCP_CLIENT_T *state = tcp_client_init();  // 初始化客户端
    if (!state) {
        return;
    }
    if (!tcp_client_open(state)) {  // 打开连接
        tcp_result(state, -1);  // 连接失败

        return;
    }
    while(!state->complete) {  // 循环等待测试完成
#if PICO_CYW43_ARCH_POLL
        // 如果使用pico_cyw43_arch_poll，需要定期调用cyw43_arch_poll
        cyw43_arch_poll();
        // 可以选择休眠直到有工作需要做
        cyw43_arch_wait_for_work_until(make_timeout_time_ms(1000));
#else
        cyw43_arch_poll();
        // 可以选择休眠直到有工作需要做
        cyw43_arch_wait_for_work_until(make_timeout_time_ms(1000));
        
        
        // 如果不使用pico_cyw43_arch_poll，可以通过中断在后台处理
        //sleep_ms(1000);  // 示例中的阻塞操作
#endif
    }
    free(state);  // 释放内存
}