/*********************************************************************************
 *      Copyright:  (C) 2025 LingYun IoT System Studio
 *                  All rights reserved.
 *
 *       Filename:  socketd_poll.c
 *    Description:  This file is TCP socket server(poll) example program.
 *
 *        Version:  1.0.0(10/20/2025)
 *         Author:  Guo Wenxue <guowenxue@gmail.com>
 *      ChangeLog:  1, Release initial version on "10/20/2025 09:48:51 PM"
 *
 ********************************************************************************/
#include <stdio.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <stdlib.h>
#include <signal.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <getopt.h>
#include <ctype.h>
#include <sys/time.h>
#include <sys/resource.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <pthread.h>
#include <poll.h>

#define BACKLOG             13
#define MAX_FDS             4096
#define ARRAY_SIZE(x)       (sizeof(x)/sizeof(x[0]))

int         g_stop = 0; /* program stop or not */

void sig_handler(int signum);
int socket_init(int port);

void print_usage(char *progname)
{
    printf("%s usage: \n", progname);
    printf("  -p(--port): sepcify listen port.\n");
    printf("  -h(--Help): print this help information.\n");

    return ;
}

int main(int argc, char **argv)
{
    int                     i, j;
    int                     rv = -1;
    int                     listen_fd,  client_fd = -1;
    int                     port = 10086;
    int                     debug = 0;
    struct sockaddr_in      cli_addr;
    socklen_t               addr_len = sizeof(cli_addr);
    char                    buf[1024];

    struct pollfd           fds_array[MAX_FDS];
    int                     need_update;
    int                     nfds = 0;
    int                     found;

    int                     ch;
    struct option           opts[] = {
        {"debug", no_argument, NULL, 'd'},
        {"port", required_argument, NULL, 'p'},
        {"help", no_argument, NULL, 'h'},
        {NULL, 0, NULL, 0}
    };

    /* 解析命令行参数 */
    while( (ch=getopt_long(argc, argv, "dp:h", opts, NULL)) != -1 )
    {
        switch(ch)
        {
            case 'd':
                debug = 1;
                break;

            case 'p':
                port=atoi(optarg);
                break;

            case 'h':
                print_usage(argv[0]);
                return 0;
        }
    }

    /* 如果不是debug调试运行，就放到后台去运行 */
    if( !debug )
        daemon(0, 0);

    /* 安装信号 */
    signal(SIGTERM, sig_handler);  /* capture kill/killall signal */
    signal(SIGINT,  sig_handler);  /* capture Ctrl+C signal */
    signal(SIGSEGV, sig_handler);  /* capture segmentation fault signal */
    signal(SIGPIPE, sig_handler);  /* capture socket error signal */

    /* 创建监听套接字 */
    if( (listen_fd = socket_init(port)) < 0 )
        return 2;

    /* 初始化 pollfd 数组 */
    for (i=0; i<ARRAY_SIZE(fds_array); i++)
    {
        fds_array[i].fd = -1;
        fds_array[i].events = 0;
        fds_array[i].revents = 0;
    }

    /* 将 listen_fd 放到 pollfd 数组的第一个元素中 */
    fds_array[0].fd = listen_fd;
    fds_array[0].events = POLLIN;
    nfds = 1;
    need_update = 0; /* 第一次进入循环只有listen_fd, 不需要更新 nfds */

    while( !g_stop )
    {
        /* 接收或关闭客户端时需要重新计算 nfds */
        if( need_update )
        {
            for(i=0, nfds=0; i<ARRAY_SIZE(fds_array); i++)
            {
                if (fds_array[i].fd>0 && i>= nfds )
                    nfds = i+1;
            }
            need_update = 0; /* 清除更新标志 */
        }

        /* poll 一直阻塞等待, 永不超时 */
        rv=poll(fds_array, nfds, -1);
        if( rv <= 0 )
        {
            /* poll()超时或有异常发生(如有信号到来)，继续下次循环 */
            printf("poll failure: %s\n", strerror(errno));
            continue;
        }

        /* 程序走到这里说明 rv>0, 有文件描述符发生事件了 */

        /* 检查是不是 listen_fd 发生了事件，如果有说明有新的客户端连接请求到来 */
        if (fds_array[0].revents & POLLIN)
        {
            /* 接收新的客户端连接请求 */
            client_fd = accept(listen_fd, (struct sockaddr*)&cli_addr, &addr_len);
            if(client_fd < 0)
            {
                printf("accept new socket failure: %s\n", strerror(errno));
                continue;
            }
            printf("Accept new client[%s:%d] with fd [%d]\n",
                    inet_ntoa(cli_addr.sin_addr), ntohs(cli_addr.sin_port), client_fd);

            /* 从符数组中找一个空位置，将新的客户端文件描述符加进去 */
            found = 0;
            for(i=0; i<ARRAY_SIZE(fds_array) ; i++)
            {
                if( fds_array[i].fd < 0 )
                {
                    fds_array[i].fd = client_fd;
                    fds_array[i].events = POLLIN;
                    found = 1;
                    need_update = 1;
                    break;
                }
            }

            /* 如果没找到说明客户端太多了，拒绝这个新的客户端连接请求 */
            if( !found )
            {
                printf("accept new client[%d] but full, so refuse it\n", client_fd);
                close(client_fd);
            }
        }

        /* fds_array[0] 是 listen_fd, 其它的则是已经连上的客户端，看是不是它们有数据到来了 */
        for(i=1; i<ARRAY_SIZE(fds_array); i++)
        {
            /* 如果位置是空的就找下一个 */
            if( fds_array[i].fd<0 )
                continue;

            /*  如果有可读事件 */
            if (fds_array[i].revents & POLLIN)
            {
                memset(buf, 0, sizeof(buf));
                if( (rv=read(fds_array[i].fd, buf, sizeof(buf))) <= 0)
                {
                    printf("socket[%d] get disconncet\n", fds_array[i].fd);
                    close(fds_array[i].fd);
                    fds_array[i].fd = -1;
                    need_update = 1;
                    continue;
                }

                printf("socket[%d] client read %d bytes data: %s\n", fds_array[i].fd, rv, buf);

                /* 将收到的字符串统一转成大写 */
                for(j=0; j<rv; j++)
                {
                    buf[i]=toupper(buf[i]);
                }

                /* 将处理后的数据返回给客户端 */
                rv=write(fds_array[i].fd, buf, rv);
                if(rv <= 0)
                {
                    printf("socket[%d] write failure: %s\n", fds_array[i].fd, strerror(errno));

                    /* 客户端关闭时要把这个位置空出来 */
                    close(fds_array[i].fd);
                    fds_array[i].fd = -1;
                    need_update = 1;
                    continue;
                }
            }

            /* 处理异常或挂起事件 */
            if (fds_array[i].revents & (POLLERR | POLLHUP | POLLNVAL))
            {
                printf("socket[%d] closed or error\n", fds_array[i].fd);
                close(fds_array[i].fd);
                fds_array[i].fd = -1;
                need_update = 1;
                continue;
            }

        }
    }

    close(listen_fd);
    return 0;
}

void sig_handler(int signum)
{
    switch(signum)
    {
        case SIGTERM:
            printf("SIGTERM signal detected(kill/killall)\n");
            break;

        case SIGSEGV:
            printf("SIGSEGV signal detected(segmentation fault)\n");
            break;

        case SIGPIPE:
            printf("SIGPIPE signal detected(socket error)\n");
            break;

        case SIGINT:
            printf("SIGINT signal detected(Ctrl+C)\n");
            break;

        default:
            break;
    }

    g_stop = 1;
}

int socket_init(int port)
{
    int                     listen_fd;
    int                     reuse = 1;
    struct rlimit           limit;
    struct sockaddr_in      serv_addr;

    /* set Max. fd limit */
    getrlimit(RLIMIT_NOFILE, &limit);
    limit.rlim_cur = limit.rlim_max;
    setrlimit(RLIMIT_NOFILE, &limit);

    /* create TCP socket */
    listen_fd = socket(AF_INET, SOCK_STREAM, 0);
    if(listen_fd < 0 )
    {
        printf("create socket failure: %s\n", strerror(errno));
        return -1;
    }

    /* resolve the "Address already in use" error when using bind() */
    setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse));

    /* bind socket server information */
    memset(&serv_addr, 0, sizeof(serv_addr));
    serv_addr.sin_family = AF_INET;
    serv_addr.sin_port = htons(port);
    serv_addr.sin_addr.s_addr = htonl(INADDR_ANY);
    if( bind(listen_fd, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) < 0 )
    {
        printf("bind socket failure: %s\n", strerror(errno));
        return -2;
    }
    printf("socket[%d] bind on port[%d] for all IP address ok\n", listen_fd, port);

    /* start listen socket */
    listen(listen_fd, BACKLOG);

    return listen_fd;
}