#include <stdio.h>
#include <string.h>
#include <winsock2.h>
#include <ws2tcpip.h>
#include <process.h>

#include "common.h"

#pragma comment(lib, "ws2_32.lib")

typedef struct {
    SOCKET sock;
    uint8_t *recv_seq;
    struct sockaddr_in6 *peer;
} recv_args;

void receiver(void *arg) {
    recv_args *a = (recv_args *)arg;
    while (1) {
        fd_set fds;
        FD_ZERO(&fds);
        FD_SET(a->sock, &fds);
        TIMEVAL tv = {0, 100000};

        if (select(a->sock + 1, &fds, NULL, NULL, &tv) > 0) {
            struct packet rcv;
            recvfrom(a->sock, (char *)&rcv, sizeof(rcv), 0, NULL, NULL);
            rcv.payload[sizeof(rcv.payload) - 1] = '\0';
            if (rcv.type == TYPE_DATA) {
                if (rcv.seq == *a->recv_seq) {
                    printf("\npeer: %s\n", rcv.payload);
                    *a->recv_seq ^= 1;
                }
                rcv.type = TYPE_ACK;
                sendto(a->sock, (char *)&rcv, 2, 0,
                       (struct sockaddr *)a->peer, sizeof(*a->peer));
            }
        }
    }
}

void send_with_stopwait(SOCKET sock,
                        struct sockaddr_in6 *peer,
                        struct packet *pkt, int pkt_len)
{
    int addr_len = sizeof(*peer);

    while (1) {
        printf("Sending seq %d\n", pkt->seq);
        sendto(sock, (char *)pkt, pkt_len, 0,
               (struct sockaddr *)peer, addr_len);

        fd_set fds;
        FD_ZERO(&fds);
        FD_SET(sock, &fds);

        TIMEVAL tv;
        tv.tv_sec = TIMEOUT_SEC;
        tv.tv_usec = 0;

        int ret = select(sock + 1, &fds, NULL, NULL, &tv);
        printf("Select returned %d\n", ret);
        if (ret > 0) {
            struct packet ack;
            recvfrom(sock, (char *)&ack, sizeof(ack), 0, NULL, NULL);
            printf("Received ack seq %d type %d\n", ack.seq, ack.type);
            if (ack.type == TYPE_ACK && ack.seq == pkt->seq) {
                printf("Correct ACK received\n");
                return;
            }
        }
        printf("[timeout] retransmit\n");
    }
}

int main(int argc, char *argv[])
{
    if (argc != 2) {
        printf("用法: %s <对方IPv6地址>\n", argv[0]);
        return 1;
    }

    WSADATA wsa;
    WSAStartup(MAKEWORD(2, 2), &wsa);

    SOCKET sock = socket(AF_INET6, SOCK_DGRAM, 0);

    struct sockaddr_in6 local = {0};
    local.sin6_family = AF_INET6;
    local.sin6_addr = in6addr_any;
    local.sin6_port = htons(PORT);
    bind(sock, (struct sockaddr *)&local, sizeof(local));

    struct sockaddr_in6 peer = {0};
    peer.sin6_family = AF_INET6;
    peer.sin6_port = htons(PORT);
    if (inet_pton(AF_INET6, argv[1], &peer.sin6_addr) != 1) {
        printf("无效的 IPv6 地址\n");
        return 1;
    }

    uint8_t send_seq = 0;
    uint8_t recv_seq = 0;
    struct packet pkt;

    memset(&pkt, 0, sizeof(pkt));

    printf("Windows UDP IPv6 Stop-Wait Chat\n");

    // 启动接收线程
    recv_args args = {sock, &recv_seq, &peer};
    _beginthread(receiver, 0, &args);

    while (1) {
        printf(">> ");
        if (!fgets(pkt.payload, sizeof(pkt.payload), stdin))
            break;

        size_t len = strlen(pkt.payload);
        if (len > 0 && pkt.payload[len - 1] == '\n') pkt.payload[len - 1] = '\0';

        pkt.seq = send_seq;
        pkt.type = TYPE_DATA;

        int pkt_len = 2 + strlen(pkt.payload) + 1;
        send_with_stopwait(sock, &peer, &pkt, pkt_len);
        send_seq ^= 1;
    }

    closesocket(sock);
    WSACleanup();
    return 0;
}