#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>
#include <spdlog/fmt/bin_to_hex.h>

#define SCG_UDP_HEAD_SIZE (11)

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));
        } else {
            SPDLOG_ERROR(TEXT(">>> Scoket In {1} Recv {0} bytes from {2}:{3} error: {4}"),
                         iRecvBytes,
                         pProxy->scgGwSock,
                         ipAddr,
                         ntohs(remoteWgAddr.sin_port),
                         WSAGetLastError());
        }

        Sleep(100);
    }
    return 0;
}

static DWORD UDPProxyRecvThread(LPVOID lpParameter) {
    bool                    isRemoteInit = false;
    sockaddr_in             localWgAddr {};
    sockaddr_in             scgAddr {};
    unsigned char           recvBuf[1500 + SCG_UDP_HEAD_SIZE];
    std::array<UINT8, 1511> arr;
    const PSCG_PROXY_INFO   pProxy  = &GetGlobalCfgInfo()->scgProxy;
    const auto              svrId   = static_cast<UINT8>(GetGlobalCfgInfo()->userCfg.cliConfig.scgTunnelAppId);
    char                   *pRecBuf = reinterpret_cast<char *>(&recvBuf[SCG_UDP_HEAD_SIZE]);

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

    // 构建 SCG UDP 包头
    recvBuf[0]  = 0x01;     // VERSION
    recvBuf[1]  = 0x09;     // Length
    recvBuf[2]  = 0xF0;     // ++++++ INFO[0] TYPE
    recvBuf[3]  = 0x04;     //        INFO[0] LENGTH
    recvBuf[4]  = 0;        //        INFO[0] VMID[0]
    recvBuf[5]  = 0;        //        INFO[0] VMID[1]
    recvBuf[6]  = 0;        //        INFO[0] VMID[2]
    recvBuf[7]  = 0;        //        INFO[0] VMID[3]
    recvBuf[8]  = 0xF1;     //        INFO[1] TYPE
    recvBuf[9]  = 0x01;     //        INFO[1] LENGTH
    recvBuf[10] = svrId;    // ------ INFO[1] SCG Service ID

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

        // 代理服务 Out
        iRecvBytes = recvfrom(pProxy->udpProxySock,
                              pRecBuf,
                              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 >= (1450 - SCG_UDP_HEAD_SIZE)) {
            SPDLOG_WARN(TEXT("!Maybe MTU overflow: Current package {0} bytes, UDP MTU 1450, SCG Head Used {1} bytes"),
                        iRecvBytes,
                        SCG_UDP_HEAD_SIZE);
        }

        if (iRecvBytes != SOCKET_ERROR) {
            int                sendBytes;
            const unsigned int id = htonl(GetGlobalCfgInfo()->curConnVmId);
            unsigned char      vmid[4];
            memcpy(vmid, &id, 4);

            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;
            }

            recvBuf[4] = vmid[0];    //        INFO[0] VMID[0]
            recvBuf[5] = vmid[1];    //        INFO[0] VMID[1]
            recvBuf[6] = vmid[2];    //        INFO[0] VMID[2]
            recvBuf[7] = vmid[3];    //        INFO[0] VMID[3]

            // 增加SCG包头数据长度
            iRecvBytes += 11;

            if (GetGlobalCfgInfo()->logLevel == spdlog::level::trace) {
                const auto start = std::begin(recvBuf);
                std::copy_n(start, iRecvBytes, std::begin(arr));
                SPDLOG_TRACE(TEXT("UDP Proxy SCG({1}/0x{2:X}) Payload: {0:Xa}"),
                             spdlog::to_hex(start, start + iRecvBytes, 16),
                             svrId,
                             id);
            }

            sendBytes = sendto(pProxy->scgGwSock,
                               reinterpret_cast<char *>(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;
}