#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) {
        WSABUF dataBuffer{};
        dataBuffer.len = BUFFER_SIZE;
        dataBuffer.buf = buffer;

        DWORD bytesReceived = 0;
        DWORD flags = 0;

        // 接收数据
        if (WSARecv(clientSocket, &dataBuffer, 1, &bytesReceived, &flags,
                    nullptr, nullptr) == SOCKET_ERROR) {
            std::cerr << "Failed to receive data: " << WSAGetLastError()
                      << std::endl;
            break;
        }

        // 检查接收到的数据大小
        if (bytesReceived > 0) {
            // 将数据写入文件
            // 将数据写入文件
            std::fwrite(buffer, sizeof(char), bytesReceived, file);
            allLen += bytesReceived;
            std::cout << "already recv: " << allLen << "\n";
        } else if (bytesReceived == 0) {
            // 客户端关闭连接
            std::cout << "Client disconnected." << std::endl;
            break;
        }
    }

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

    return 0;
}