#include "pch.h"

#include <strsafe.h>
#include <spdlog/spdlog.h>
#include "globalcfg.h"
#include "usrerr.h"
#include "misc.h"

#include <winsock2.h>
#include <ws2tcpip.h>

typedef struct {
    bool   isRunning;
    HANDLE proxyThread;
    HANDLE proxyHandle;
    bool   isExitSvr;
} UDP_PROXY_CTX, *PUDP_PROXY_CTX;

typedef struct {
    TCHAR         streamFilter[1024];
    TCHAR         targetIp[MAX_IP_LEN];
    UINT16        targetPort;
    int           vmId;
    int           svrId;
    UDP_PROXY_CTX udpCtx;
} PROXY_INFO, *PPROXY_INFO;

void StopUDPProxyServer() {
    const PSCG_PROXY_INFO pProxy = &GetGlobalCfgInfo()->scgProxy;
    pProxy->exitNow              = true;

    if (pProxy->hProxyTunnelThread) {
        if (WaitForSingleObject(pProxy->hProxyTunnelThread, 10 * 1000) == WAIT_TIMEOUT) {
            SPDLOG_ERROR(TEXT("Waitting HTTP Service clost timeout"));
        }

        closesocket(pProxy->udpProxySock);
    }

    if (pProxy->hProxySCGThread) {
        if (WaitForSingleObject(pProxy->hProxySCGThread, 10 * 1000) == WAIT_TIMEOUT) {
            SPDLOG_ERROR(TEXT("Waitting HTTP Service clost timeout"));
        }

        closesocket(pProxy->scgGwSock);
    }
}

static DWORD UDPProxvRemoteThread(LPVOID lpParameter) {
    const auto            pPeerSock = static_cast<sockaddr_in *>(lpParameter);
    const PSCG_PROXY_INFO pProxy    = &GetGlobalCfgInfo()->scgProxy;

    while (pPeerSock && !pProxy->exitNow) {
        sockaddr_in remoteWgAddr {};
        TCHAR       ipAddr[MAX_IP_LEN];

        int  addrSize = sizeof(SOCKADDR);
        char recvBuf[1500];
        int  iRecvBytes;

        // 代理服务 In
        iRecvBytes = recvfrom(pProxy->scgGwSock,
                              recvBuf,
                              1500,
                              0,
                              reinterpret_cast<SOCKADDR *>(&remoteWgAddr),
                              &addrSize);

        memset(ipAddr, 0, MAX_IP_LEN);
        InetNtop(AF_INET, &remoteWgAddr.sin_addr.s_addr, ipAddr, MAX_IP_LEN);
        SPDLOG_TRACE(TEXT(">>> Scoket In {1} Recv {0} bytes from {2}:{3}"),
                     iRecvBytes,
                     pProxy->scgGwSock,
                     ipAddr,
                     ntohs(remoteWgAddr.sin_port));

        if (iRecvBytes != SOCKET_ERROR) {
            int sendBytes = sendto(pProxy->udpProxySock,
                                   recvBuf,
                                   iRecvBytes,
                                   0,
                                   reinterpret_cast<SOCKADDR *>(pPeerSock),
                                   sizeof(SOCKADDR));
            memset(ipAddr, 0, MAX_IP_LEN);
            InetNtop(AF_INET, &pPeerSock->sin_addr.s_addr, ipAddr, MAX_IP_LEN);
            SPDLOG_TRACE(TEXT("<<< Scoket In  Send {0} bytes to {2}:{3}"),
                         sendBytes,
                         pProxy->udpProxySock,
                         ipAddr,
                         ntohs(pPeerSock->sin_port));
        }

        Sleep(100);
    }
    return 0;
}

static DWORD UDPProxyRecvThread(LPVOID lpParameter) {
    bool                  isRemoteInit = false;
    sockaddr_in           localWgAddr {};
    sockaddr_in           scgAddr {};
    const PSCG_PROXY_INFO pProxy = &GetGlobalCfgInfo()->scgProxy;
    pProxy->exitNow              = false;

    scgAddr.sin_family = AF_INET;
    scgAddr.sin_port   = htons(pProxy->scgGwPort);
    InetPton(AF_INET, pProxy->scgIpAddr, &scgAddr.sin_addr.s_addr);

    while (!pProxy->exitNow) {
        TCHAR ipAddr[MAX_IP_LEN];
        int   addrSize = sizeof(SOCKADDR);
        char  recvBuf[1500];
        int   iRecvBytes;

        // 代理服务 Out
        iRecvBytes = recvfrom(pProxy->udpProxySock,
                              recvBuf,
                              1500,
                              0,
                              reinterpret_cast<SOCKADDR *>(&localWgAddr),
                              &addrSize);

        InetNtop(AF_INET, &localWgAddr.sin_addr.s_addr, ipAddr, MAX_IP_LEN);
        SPDLOG_TRACE(TEXT(">>> Scoket Out {1} Recv {0} bytes from {2}:{3}"),
                     iRecvBytes,
                     pProxy->udpProxySock,
                     ipAddr,
                     ntohs(localWgAddr.sin_port));

        if (iRecvBytes != SOCKET_ERROR) {
            int sendBytes;

            if (!isRemoteInit) {
                HANDLE handle;
                isRemoteInit = true;
                // 创建远端接收线程
                handle       = CreateThread(nullptr,                 // Thread attributes
                                      0,                       // Stack size (0 = use default)
                                      UDPProxvRemoteThread,    // Thread start address
                                      &localWgAddr,            // Parameter to pass to the thread
                                      0,                       // Creation flags
                                      nullptr);                // Thread id

                if (handle == nullptr) {
                    SPDLOG_ERROR("Create Thread failed with error = {0}", GetLastError());
                    closesocket(pProxy->udpProxySock);
                    closesocket(pProxy->scgGwSock);
                    pProxy->exitNow = true;
                    return -ERR_CREATE_THREAD;
                }

                pProxy->hProxySCGThread = handle;
            }

            sendBytes = sendto(pProxy->scgGwSock,
                               recvBuf,
                               iRecvBytes,
                               0,
                               reinterpret_cast<SOCKADDR *>(&scgAddr),
                               sizeof(SOCKADDR));
            memset(ipAddr, 0, MAX_IP_LEN);
            InetNtop(AF_INET, &scgAddr.sin_addr.s_addr, ipAddr, MAX_IP_LEN);
            SPDLOG_TRACE(TEXT("<<< Scoket Out  Send {0} bytes to {2}:{3}"),
                         sendBytes,
                         pProxy->scgGwSock,
                         ipAddr,
                         ntohs(scgAddr.sin_port));
        }

        Sleep(100);
    }
    return 0;
}

int CreateUDPProxyServer() {
    HANDLE                handle;
    int                   ret;
    int                   addrSize = sizeof(sockaddr_in);
    sockaddr_in           server {};
    sockaddr_in           bindAddr {};
    SOCKET                sock;
    const PSCG_PROXY_INFO pProxy = &GetGlobalCfgInfo()->scgProxy;

    // 创建本地 SOCKET 代理服务器
    sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sock == INVALID_SOCKET) {
        SPDLOG_ERROR("Create UDP Socket failed with error = {0}", WSAGetLastError());
        return -ERR_SOCKET_CREATE;
    }

    server.sin_family = AF_INET;
    server.sin_addr.s_addr = htonl(INADDR_ANY);
    server.sin_port   = htons(0);

    if (bind(sock, reinterpret_cast<SOCKADDR *>(&server), sizeof(SOCKADDR)) == SOCKET_ERROR) {
        closesocket(sock);
        SPDLOG_ERROR("Bind local UDP Socket failed with error = {0}", WSAGetLastError());
        return -ERR_SOCKET_BIND;
    }

    if ((ret = getsockname(sock, reinterpret_cast<SOCKADDR *>(&bindAddr), &addrSize)) != 0) {
        closesocket(sock);
        SPDLOG_ERROR("Get UDP Socket bind port failed with error = {0}, {1}", WSAGetLastError(), ret);
        return -ERR_SOCKET_BIND;
    }

    // 保存 UDP 代理服务器信息
    pProxy->udpProxySock = sock;
    pProxy->proxyPort    = ntohs(bindAddr.sin_port);

    SPDLOG_DEBUG(TEXT("Proxy Server socket {0} bind {1} prot"), sock, pProxy->proxyPort);

    // 创建SCG SOCKET 连接客户端服务
    sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sock == INVALID_SOCKET) {
        SPDLOG_ERROR("Create UDP Socket failed with error = {0}", WSAGetLastError());
        closesocket(pProxy->udpProxySock);
        return -ERR_SOCKET_CREATE;
    }

    pProxy->scgGwSock = sock;

    // 创建代理服务发送线程
    handle = CreateThread(nullptr,               // Thread attributes
                          0,                     // Stack size (0 = use default)
                          UDPProxyRecvThread,    // Thread start address
                          nullptr,               // Parameter to pass to the thread
                          0,                     // Creation flags
                          nullptr);              // Thread id

    if (handle == nullptr) {
        SPDLOG_ERROR("Create Thread failed with error = {0}", GetLastError());
        closesocket(pProxy->udpProxySock);
        closesocket(pProxy->scgGwSock);
        return -ERR_CREATE_THREAD;
    }

    pProxy->hProxyTunnelThread = handle;    

    return ERR_SUCCESS;
}