/*

 * Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved.

 * virtCCA_sdk is licensed under the Mulan PSL v2.

 * You can use this software according to the terms and conditions of the Mulan PSL v2.

 * You may obtain a copy of Mulan PSL v2 at:

 *     http://license.coscl.org.cn/MulanPSL2

 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR

 * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR

 * PURPOSE.

 * See the Mulan PSL v2 for more details.

 */

#include <sys/socket.h>

#include <netinet/in.h>

#include <arpa/inet.h>

#include <unistd.h>

#include <cstring>

#include <cstdlib>

#include <string>

#include <vector>

#include <memory>

#include <iostream>

#include <netinet/tcp.h>

#include <future>

#include <fstream>

#include <unistd.h>

#include <fcntl.h>

#include <sys/types.h>

#include <sys/wait.h>

#include "data_guard_mpc.h"



#define PORT 8080

#define BUFFER_SIZE 1024

#define BACKLOG_SIZE 16  // 请求队列最大长度

#define SLEEP_FIVE 5

#define NODES_SIZE 2



int g_sendSockfd = -1;

int g_recvSockfd = -1;



int CallMpcTee(int nodeId, std::string inputFileName, std::string sendIp, std::string recvIp, short port1, short port2);

int senddata(struct TeeNodeInfo *nodeInfo, unsigned char *buf, u64 len);

int recvdata(struct TeeNodeInfo *nodeInfo, unsigned char *buf, u64 *len);

void CloseSocket();

int InitServer(int port);

int InitClient(const char *ip, int port);



void BuildDgString(std::vector<std::string> &strings, DG_String **dg, unsigned int &size)

{

    size = strings.size();

    DG_String *dgString = new DG_String[strings.size()];

    for (size_t i = 0; i < strings.size(); i++) {

        dgString[i].str = strdup(strings[i].c_str());

        dgString[i].size = strings[i].size() + 1;

    }

    *dg = dgString;

}



void ReleaseDgString(DG_String *dg, unsigned int size)

{

    for (size_t i = 0; i < size; i++) {

        free(dg[i].str);

    }

    delete[] dg;

}



TeeNodeInfos BuildTeeNodeInfos(short port1, short port2)

{

    struct TeeNodeInfo teeNodes[NODES_SIZE];

    teeNodes[0].nodeId = 0;

    teeNodes[1].nodeId = 1;

    struct TeeNodeInfos allNodes;

    allNodes.nodeInfo = teeNodes;

    allNodes.size = NODES_SIZE;

    return allNodes;

}



int main(int argc, char **argv)

{

    if (argc < 7) {

        printf("params: send port, recv port, current nodeId, inputfileName, firstserver,\n");

        exit(-1);

    }



    int sendPort = (int) std::atoi(argv[1]);

    int recvPort = (int) std::atoi(argv[2]);

    int nodeId = std::atoi(argv[3]);

    std::string inputFileName = argv[4];

    TeeNodeInfos teeNodeInfo{};

    // firstSer 谁先启动，0表示先启动

    int firstSer = std::atoi(argv[5]);

        // 发送数据的Ip

    std::string sendIp = argv[6];

    // 接收数据的Ip

    std::string recvIp = argv[7];



    if (firstSer == 0) {

        std::future<int> recvFuture = std::async(InitServer, sendPort);

        g_recvSockfd = recvFuture.get();

        std::future<int> sendFuture = std::async(InitClient, recvIp.data(), recvPort);

        g_sendSockfd = sendFuture.get();

        sleep(SLEEP_FIVE);

        CallMpcTee(nodeId, inputFileName, sendIp, recvIp, sendPort, recvPort);

    } else {

        std::future<int> sendFuture = std::async(InitClient, recvIp.data(), recvPort);

        g_sendSockfd = sendFuture.get();

        std::future<int> recvFuture = std::async(InitServer, sendPort);

        g_recvSockfd = recvFuture.get();

        sleep(SLEEP_FIVE);

        CallMpcTee(nodeId, inputFileName, recvIp, sendIp, recvPort, sendPort);

    }



    sleep(SLEEP_FIVE);

    CloseSocket();

    return 0;

}



static int64_t GetCurrentTimestampMs()

{

    auto now = std::chrono::system_clock::now();

    auto duration = std::chrono::duration_cast<std::chrono::microseconds>(now.time_since_epoch());

    return duration.count() / 1000;

}





int CallMpcTee(int nodeId, std::string inputFileName, std::string sendIp, std::string recvIp, short port1, short port2)

{

    printf("---------DG_InitConfigOpts------------\n");

    DG_ConfigOpts *opts = nullptr;

    int rv = DG_InitConfigOpts(DG_BusinessType::MPC, &opts);

    if (rv != 0) {

        printf("DG_InitConfigOpts error!code:%d\n", rv);

        return rv;

    }



    printf("---------Build_Config------------\n");

    void *teeCfg = NULL;

    rv = opts->init(&teeCfg);

    if (rv != 0) {

        printf(" opts->init(dgCfg)!-%d\n", rv);

        return rv;

    }

    opts->setIntValue(teeCfg, DG_CON_MPC_TEE_INT_NODEID, nodeId);

    opts->setIntValue(teeCfg, DG_CON_MPC_TEE_INT_FXP_BITS, 8);

    opts->setIntValue(teeCfg, DG_CON_MPC_TEE_INT_THREAD_COUNT, 16);



    TEE_NET_RES teeNet = {senddata, recvdata};

    DG_Void netFunc;

    netFunc.data = &teeNet;

    netFunc.size = sizeof(TEE_NET_RES);

    opts->setVoidValue(teeCfg, DG_CON_MPC_TEE_VOID_NET_API, &netFunc);



    printf("---------DG_InitTeeMpcOpts------------\n");

    DG_PrivateSet_Opts teeOpts = DG_InitPsiOpts();



    printf("---------initTeeMpcSql------------\n");

    struct DG_TeeCtx *dgTee = nullptr;

    rv = teeOpts.initTeeCtx(teeCfg, &dgTee);

    if (rv != 0) {

        printf("tee init error.-%d\n", rv);

        return rv;

    }

    printf("---------setTeeNodeInfos------------\n");

    struct TeeNodeInfo teeNodes[NODES_SIZE];

    teeNodes[0].nodeId = 0;

    teeNodes[1].nodeId = 1;

    struct TeeNodeInfos allNodes;

    allNodes.nodeInfo = teeNodes;

    allNodes.size = NODES_SIZE;

    rv = teeOpts.setTeeNodeInfos(dgTee, &allNodes);

    if (rv != 0) {

        printf("tee set node info error.-%d\n", rv);

        return rv;

    }

    printf("---------executePrivateSetOpts------------\n");

    std::ifstream ifStream(inputFileName);

    std::string line;

    std::vector<std::string> datas;

    while (std::getline(ifStream, line)) {

        datas.push_back(line);

    }

    unsigned int size;

    DG_String *strings = nullptr;

    BuildDgString(datas, &strings, size);

    DG_TeeInput teeInput;

    teeInput.data.strings = strings;

    teeInput.size = size;

    teeInput.dataType = MPC_STRING;

    printf("++++++++++input data size:%lu\n", teeInput.size);



    DG_TeeOutput *output = nullptr;

    uint64_t start = GetCurrentTimestampMs();

    int res = teeOpts.calculate(dgTee, PSI, &teeInput, &output, TEE_OUTPUT_INDEX);

    printf("total time:%lu\n", GetCurrentTimestampMs() - start);

    if (res != 0 || output == nullptr) {

        printf("calc psi result:%d\n", res);

        return res;

    }



    int randomNumber = std::rand() % 100 + 10; // 生成 [1, 100] 范围内的随机数

    auto now = std::chrono::system_clock::now();

    // 将时间转化为自1970年1月1日以来的毫秒数

    auto duration = now.time_since_epoch();

    auto timestamp = std::chrono::duration_cast<std::chrono::milliseconds>(duration).count();

    std::string outFile = "output-" + std::to_string(randomNumber) + std::to_string(timestamp) + \

        "-" + (nodeId == 0 ? "server" : "client") + ".csv";

    int outFp = open(outFile.c_str(), O_RDWR | O_CREAT);

    printf("++++++++++++++output:%lu\n", output->size);

    if (output->dataType == MPC_STRING) {

        for (int j = 0; j < output->size; j++) {

            std::string conent = std::string(output->data.strings[j].str) + "\n";

            write(outFp, conent.c_str(), conent.size());

        }

    } else {

        for (int j = 0; j < output->size; j++) {

            std::string conent = std::to_string(output->data.u64Numbers[j]) + "\n";

            write(outFp, conent.c_str(), conent.size());

        }

    }



    teeOpts.releaseOutput(&output); // 释放输出

    printf("release output:%d\n", res);

    ReleaseDgString(strings, size);  // 释放输入

    teeOpts.releaseTeeCtx(&dgTee); // 释放dgTee

    opts->release(&teeCfg);  // 释放config teeCfg

    DG_ReleaseConfigOpts(&opts); // 释放dataguard操作集opts

    close(outFp);

    return 0;

}



int InitServer(int port)

{

    int serverFd = -1;

    int clientFd = -1;

    sockaddr_in serverAddr{};

    socklen_t addrLen = sizeof(serverAddr);



    // 创建服务器套接字

    serverFd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);

    if (serverFd == -1) {

        std::cerr << "Failed to create socket: " << strerror(errno)

                  << std::endl;

        return -1;

    }

    std::cout<<"port1 = "<<port<<std::endl;

    // 设置服务器地址

    serverAddr.sin_family = AF_INET;

    serverAddr.sin_addr.s_addr = INADDR_ANY;

    serverAddr.sin_port = htons(port);

    std::cout<<"port2 = "<<port<<std::endl;

    // 绑定套接字

    int on = 1;

    setsockopt(serverFd, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(int));

    if (bind(serverFd, (sockaddr *)&serverAddr, sizeof(serverAddr)) == -1) {

        std::cerr << "Failed to bind socket: " << strerror(errno) << std::endl;

        close(serverFd);

        return -1;

    }



    // 监听套接字

    if (listen(serverFd, BACKLOG_SIZE) == -1) {

        std::cerr << "Failed to listen on socket: " << strerror(errno)

                  << std::endl;

        close(serverFd);

        return -1;

    }



    std::cout << "Server started. Listening on port " << port << "..."

              << std::endl;



    // 接受一个连接

    clientFd = accept(serverFd, (sockaddr *)&serverAddr, &addrLen);

    if (clientFd == -1) {

        std::cerr << "Failed to accept client connection: " << strerror(errno)

                  << std::endl;

        close(serverFd);

        return -1;

    }

    int opt = 1;

    setsockopt(clientFd, IPPROTO_TCP, TCP_QUICKACK, (void *)&opt, sizeof(int));

    setsockopt(clientFd, IPPROTO_TCP, TCP_NODELAY, (void *)&opt, sizeof(int));



    // 关闭服务器套接字，因为我们只需要处理一个连接

    close(serverFd);



    std::cout << "New connection from " << inet_ntoa(serverAddr.sin_addr)

              << ":" << ntohs(serverAddr.sin_port) << std::endl;



    return clientFd;

}



int InitClient(const char *ip, int port)

{

    int servFd = -1;

    struct sockaddr_in servAddr{};



    // 创建 Socket 文件描述符

    servFd = socket(AF_INET, SOCK_STREAM, 0);

    if (servFd < 0) {

        std::cerr << "Socket creation error: " << strerror(errno) << std::endl;

        return -1;

    }



    // 设置 Server 连接信息

    servAddr.sin_family = AF_INET;

    servAddr.sin_port = htons(port);

    // 将 ip 地址从文本转换为二进制格式

    if (inet_pton(AF_INET, ip, &servAddr.sin_addr) <= 0) {

        std::cerr << "Invalid address/ Address not supported" << std::endl;

        close(servFd);

        return -1;

    }



    while (true) {

        if (connect(servFd, (struct sockaddr *)&servAddr,

                    sizeof(servAddr)) == -1) {

            std::cerr << "Connection to server failed: " << strerror(errno)

                      << std::endl;

            sleep(1);

            continue;

        }

        std::cout << "Connection to server successful" << std::endl;

        // 设置套接字选项

        int opt = 1;

        if (setsockopt(servFd, IPPROTO_TCP, TCP_QUICKACK, (void *)&opt,

                       sizeof(int)) == -1) {

            std::cerr << "Failed to set TCP_QUICKACK option" << std::endl;

            close(servFd);

            return -1;

        }

        if (setsockopt(servFd, IPPROTO_TCP, TCP_NODELAY, (void *)&opt,

                       sizeof(int)) == -1) {

            std::cerr << "Failed to set TCP_NODELAY option" << std::endl;

            close(servFd);

            return -1;

        }

        break;

    }



    return servFd;

}



int senddata(struct TeeNodeInfo *nodeInfo, unsigned char *buf, u64 len)

{

    return send(g_sendSockfd, buf, len, 0);

}



int recvdata(struct TeeNodeInfo *nodeInfo, unsigned char *buf, u64 *len)

{

    ssize_t nread;

    ssize_t tmpReadLen = 0;

    uint64_t remainReadLen = *len;



    while (remainReadLen > 0) {

        nread = read(g_recvSockfd, buf + tmpReadLen, *len - tmpReadLen);

        if (nread == 0) {

            printf("socket shut down\n");

            return 0;

        }

        if (nread < 0) {

            if (EINTR == errno || EAGAIN == errno) {

                continue;

            } else {

                printf("[%s] socket error\n", strerror(errno));

                return nread;

            }

        }

        tmpReadLen += nread;

        remainReadLen -= nread;

    }

    return *len;

}



void CloseSocket()

{

    close(g_sendSockfd);

    g_sendSockfd = NULL;

    close(g_recvSockfd);

    g_recvSockfd = NULL;

}