#include "SocketComm.h"

SocketComm::SocketComm() {
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        std::cerr << "WSAStartup failed\n";
    }
}
SocketComm::~SocketComm() {
    closeSocket();
    WSACleanup();
}

bool SocketComm::setNonBlocking(SOCKET s, bool nb) {
    u_long mode = nb ? 1UL : 0UL;
    return ioctlsocket(s, FIONBIO, &mode) == 0;
}

bool SocketComm::createClientSocket(const std::string& serverIP, uint16_t serverPort) {
    sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (sock == INVALID_SOCKET) return false;

    sockaddr_in addr{}; addr.sin_family = AF_INET; addr.sin_port = htons(serverPort);
    if (InetPton(AF_INET, serverIP.c_str(), &addr.sin_addr) <= 0) return false;

    if (connect(sock, (sockaddr*)&addr, sizeof(addr)) == SOCKET_ERROR) {
        std::cerr << "connect failed: " << WSAGetLastError() << "\n";
        closesocket(sock); sock = INVALID_SOCKET; return false;
    }
    return true;
}

bool SocketComm::createServerSocket(uint16_t port) {
    listenSock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (listenSock == INVALID_SOCKET) return false;

    // ¸´ÓõØÖ·
    BOOL yes = 1; setsockopt(listenSock, SOL_SOCKET, SO_REUSEADDR, (char*)&yes, sizeof(yes));

    sockaddr_in addr{}; addr.sin_family = AF_INET; addr.sin_addr.s_addr = INADDR_ANY; addr.sin_port = htons(port);
    if (bind(listenSock, (sockaddr*)&addr, sizeof(addr)) == SOCKET_ERROR) return false;
    if (listen(listenSock, SOMAXCONN) == SOCKET_ERROR) return false;

    // ·Ç×èÈû¼àÌý
    setNonBlocking(listenSock, true);
    std::cout << "Server listening on port " << port << "\n";
    return true;
}

bool SocketComm::acceptOne(SOCKET& outClient, std::string& outIp, uint16_t& outPort) {
    sockaddr_in caddr{}; int len = sizeof(caddr);
    SOCKET cs = accept(listenSock, (sockaddr*)&caddr, &len);
    if (cs == INVALID_SOCKET) {
        int e = WSAGetLastError();
        if (e == WSAEWOULDBLOCK) return false; // µ±Ç°Ã»ÓÐÐÂÁ¬½Ó
        std::cerr << "accept error: " << e << "\n"; return false;
    }
    setNonBlocking(cs, true);
    char ipbuf[INET_ADDRSTRLEN]{};
    InetNtop(AF_INET, &caddr.sin_addr, ipbuf, INET_ADDRSTRLEN);
    outIp = ipbuf; outPort = ntohs(caddr.sin_port);
    outClient = cs;
    return true;
}

bool SocketComm::recvFrom(SOCKET s, std::vector<uint8_t>& buffer, bool& peerClosed) {
    peerClosed = false;
    char tmp[4096];
    int r = recv(s, tmp, sizeof(tmp), 0);
    if (r > 0) {
        buffer.assign(tmp, tmp + r);
        return true;
    }
    if (r == 0) { // ¶Ô¶ËÕý³£¹Ø±Õ
        peerClosed = true;
        return false;
    }
    int e = WSAGetLastError();
    if (e == WSAEWOULDBLOCK) return false; // ÔÝÎÞÊý¾Ý
    // ÆäËü´íÎóÒ²ÈÏΪÁ¬½Ó²»¿ÉÓÃÁË
    peerClosed = true;
    return false;
}

bool SocketComm::sendTo(SOCKET s, const std::vector<uint8_t>& data) {
    int sent = send(s, reinterpret_cast<const char*>(data.data()), (int)data.size(), 0);
    return sent == (int)data.size();
}

void SocketComm::closeClient(SOCKET s) {
    if (s != INVALID_SOCKET) closesocket(s);
}

void SocketComm::closeSocket() {
    if (listenSock != INVALID_SOCKET) { closesocket(listenSock); listenSock = INVALID_SOCKET; }
    if (sock != INVALID_SOCKET) { closesocket(sock);       sock = INVALID_SOCKET; }
}

bool SocketComm::sendDataSingle(const std::vector<uint8_t>& data) {
    if (sock == INVALID_SOCKET) return false;
    return sendTo(sock, data);
}
bool SocketComm::recvSingle(std::vector<uint8_t>& buffer) {
    if (sock == INVALID_SOCKET) return false;
    return recvFrom(sock, buffer);
}