#include <iostream>
#include <thread>
#include <vector>
#include <map>
#include <mutex>
#include <atomic>
#include <sys/socket.h>
#include <netinet/in.h>
#include <unistd.h>
#include <arpa/inet.h>

#define BUFFER_SIZE 1024

std::vector<std::thread> threads; // 存储所有连接线程
std::map<int, int> sockets; // 存储套接字
std::mutex sockets_mutex; // 用于同步访问 sockets
std::atomic<bool> exit_flag(false); // 原子退出标志


// 维护连接的函数
void KeepConnection(int sock) {
    char buffer[BUFFER_SIZE];
    int nbytes;
    while (!exit_flag) {
        nbytes = recv(sock, buffer, BUFFER_SIZE, 0);
        if (nbytes <= 0) {
            std::cout << "Gateway disconnected" << std::endl;
            break;
        }
        std::cout << "Received from gateway (Socket " << sock << "): " << buffer << std::endl;
    }

    {
        std::lock_guard<std::mutex> lock(sockets_mutex);
        sockets.erase(sock); // 从集合中移除套接字
    }
    close(sock);
}
// 连接到网关的函数
void ConnectToGateway(const std::string& ip, int port) {
    int sock = socket(AF_INET, SOCK_STREAM, 0);
    if (sock < 0) {
        perror("Could not create socket");
        return;
    }

    struct sockaddr_in server_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(port);
    if (inet_pton(AF_INET, ip.c_str(), &server_addr.sin_addr) <= 0) {
        perror("Invalid address");
        close(sock);
        return;
    }

    if (connect(sock, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("Connection to the server failed");
        close(sock);
        return;
    }

    {
        std::lock_guard<std::mutex> lock(sockets_mutex);
        sockets[sock] = 1; // 将套接字添加到集合中
    }

    std::cout << "Connected to gateway at " << ip << ":" << port << " (Socket " << sock << ")" << std::endl;

    // 在新线程中维护连接
    threads.emplace_back(KeepConnection, sock);
}



int main() {
    std::string command, ip;
    int port;
    bool in_connect_mode = false;

    while (true) {
        std::cout << "Enter command (connect/exitconnect/exit): ";
        std::cin >> command;

        if (command == "connect") {
            if (in_connect_mode) {
                std::cout << "Already in connect mode." << std::endl;
                continue;
            }
            in_connect_mode = true; // 进入连接模式
            std::cout << "Enter the gateway IP address: ";
            std::cin >> ip;
            std::cout << "Enter the gateway port: ";
            std::cin >> port;
            ConnectToGateway(ip, port);
        } else if (command == "exitconnect") {
            if (!in_connect_mode) {
                std::cout << "Not in connect mode." << std::endl;
                continue;
            }
            in_connect_mode = false; // 退出连接模式
            std::cout << "Exiting connect mode." << std::endl;
        } else if (command == "exit") {
            exit_flag = true; // 设置退出标志
            std::cout << "Exiting program." << std::endl;
            // 关闭所有套接字
            for (auto& sock_pair : sockets) {
                shutdown(sock_pair.first, SHUT_RDWR); // 关闭套接字的发送和接收
                close(sock_pair.first);
            }
            sockets.clear();
            // 等待所有线程结束
            for (auto& thread : threads) {
                if (thread.joinable()) {
                    thread.join();
                }
            }
            threads.clear();
            break; // 退出主循环
        } else {
            std::cout << "Unknown command" << std::endl;
        }
    }

    return 0;
}