#include <WS2tcpip.h>
#include <WinSock2.h>

#include <cstdio>
#include <fstream>
#include <iostream>
#include <string>


#define BUFFER_SIZE 1024

int main(int argc, char** argv)
{
    if (argc == 1) {
        std::cout << "Need Param, Example: server 127.0.0.1 9696 E:\\2.bmp\n";
        return 0;
    }

    // 初始化 Winsock
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        std::cerr << "Failed to initialize Winsock." << std::endl;
        return 1;
    }

    // 创建套接字
    SOCKET serverSocket = socket(AF_INET, SOCK_STREAM, 0);
    if (serverSocket == INVALID_SOCKET) {
        std::cerr << "Failed to create socket: " << WSAGetLastError()
                  << std::endl;
        WSACleanup();
        return 1;
    }

    // 绑定地址和端口
    sockaddr_in serverAddr{};
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_port = htons(std::stoi(std::string(argv[2])));
    serverAddr.sin_addr.s_addr = inet_addr(argv[1]);
    ;
    if (bind(serverSocket, reinterpret_cast<sockaddr*>(&serverAddr),
             sizeof(serverAddr)) == SOCKET_ERROR) {
        std::cerr << "Failed to bind socket: " << WSAGetLastError()
                  << std::endl;
        closesocket(serverSocket);
        WSACleanup();
        return 1;
    }

    // 监听连接请求
    if (listen(serverSocket, 1) == SOCKET_ERROR) {
        std::cerr << "Failed to listen on socket: " << WSAGetLastError()
                  << std::endl;
        closesocket(serverSocket);
        WSACleanup();
        return 1;
    }

    std::cout << "server ip: " << argv[1] << " port: " << argv[2] << std::endl;
    std::cout << "Server started. Waiting for client connection..."
              << std::endl;

    // 接受客户端连接
    SOCKET clientSocket = accept(serverSocket, nullptr, nullptr);
    if (clientSocket == INVALID_SOCKET) {
        std::cerr << "Failed to accept client connection: " << WSAGetLastError()
                  << std::endl;
        closesocket(serverSocket);
        WSACleanup();
        return 1;
    }

    std::cout << "Client connected. Receiving data..." << std::endl;

    // 创建接收缓冲区
    char   buffer[BUFFER_SIZE];
    size_t allLen = 0;
    // 打开文件以保存接收的数据
    FILE* file = std::fopen(argv[3], "wb");
    if (!file) {
        std::cerr << "Failed to open file for writing." << std::endl;
        closesocket(clientSocket);
        closesocket(serverSocket);
        WSACleanup();
        return 1;
    }

    // 循环接收数据
    while (true) {
        // 创建事件对象
        WSAEVENT event = WSACreateEvent();
        if (event == WSA_INVALID_EVENT) {
            std::cerr << "Failed to create event: " << WSAGetLastError()
                      << std::endl;
            break;
        }

        // 将事件对象与套接字关联
        if (WSAEventSelect(clientSocket, event, FD_READ) == SOCKET_ERROR) {
            std::cerr << "Failed to associate event with socket: "
                      << WSAGetLastError() << std::endl;
            WSACloseEvent(event);
            break;
        }

        // 等待事件对象变为 signaled 状态
        DWORD waitResult =
            WSAWaitForMultipleEvents(1, &event, FALSE, WSA_INFINITE, FALSE);
        if (waitResult == WSA_WAIT_FAILED) {
            std::cerr << "Failed to wait for event: " << WSAGetLastError()
                      << std::endl;
            WSACloseEvent(event);
            break;
        }

        // 检查事件对象状态
        if (waitResult == WSA_WAIT_EVENT_0) {
            WSANETWORKEVENTS networkEvents;
            if (WSAEnumNetworkEvents(clientSocket, event, &networkEvents) ==
                SOCKET_ERROR) {
                std::cerr << "Failed to enumerate network events: "
                          << WSAGetLastError() << std::endl;
                WSACloseEvent(event);
                break;
            }

            // 检查是否有数据可读
            if (networkEvents.lNetworkEvents & FD_READ) {
                if (networkEvents.iErrorCode[FD_READ_BIT] != 0) {
                    std::cerr << "Error in FD_READ: "
                              << networkEvents.iErrorCode[FD_READ_BIT]
                              << std::endl;
                    WSACloseEvent(event);
                    break;
                }

                // 接收数据
                int bytesRead = recv(clientSocket, buffer, BUFFER_SIZE, 0);
                if (bytesRead > 0) {
                    // 将数据写入文件
                    std::fwrite(buffer, sizeof(char), bytesRead, file);
                    allLen += bytesRead;
                    std::cout << "already recv: " << allLen << "\n";
                } else if (bytesRead == 0) {
                    // 客户端关闭连接
                    std::cout << "Client disconnected." << std::endl;
                    WSACloseEvent(event);
                    break;
                } else {
                    // 发生错误
                    std::cerr << "Failed to receive data: " << WSAGetLastError()
                              << std::endl;
                    WSACloseEvent(event);
                    break;
                }
            }
        }

        WSACloseEvent(event);
    }

    // 关闭套接字和清理资源
    closesocket(clientSocket);
    closesocket(serverSocket);
    WSACleanup();
    std::fclose(file);

    std::cout << "Done .... file: " << argv[3] << std::endl;
    return 0;
}