/*
 * Copyright (c) 2023-2024 Ian Marco Moffett and the Osmora Team.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 * 3. Neither the name of Hyra nor the names of its
 *    contributors may be used to endorse or promote products derived from
 *    this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include <sys/socket.h>
#include <netinet/in.h>
#include <ostp/net/auth.h>
#include <ostp/net/stpsession.h>
#include <arpa/inet.h>
#include <ostp/server.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <stdio.h>
#include <pthread.h>

#define MAX_BACKLOG 4
#define LISTEN_PORT 5352

struct recv_args {
    void *buf;
    struct ostp_client *c;
    struct ostp_listener *lp;
    size_t n;
    pthread_t td;
};

static int
handle_client(struct sockaddr_in *caddr, struct ostp_client *c, struct ostp_listener *lp)
{
    struct session_request srq;
    ssize_t nread;

    /* Try to read in the session request */
    if ((nread = read(c->sockfd, &srq, sizeof(srq))) < 0) {
        printf("Read failure...\n");
        listener_close(lp, c);
        return -1;
    }

    if (nread == 0) {
        printf("Connection closed by peer\n");
        listener_close(lp, c);
        return -1;
    }

    /* Is this even a session request? */
    if (nread != sizeof(srq)) {
        printf("Rejecting data - not a session request...\n");
        listener_close(lp, c);
        return -1;
    }

    /* Handle the session request */
    if (handle_srq(c, lp, &srq) < 0) {
        listener_close(lp, c);
        return -1;
    }

    return 0;
}

/*
 * When we get data from a client, a thread is sent
 * here to handle it.
 */
static void *
recv_td(void *args)
{
    struct recv_args *ap = args;
    struct ostp_listener *lp;

    lp = ap->lp;
    if (lp->on_recv != NULL) {
        lp->on_recv(ap->c, ap->buf, ap->n);
    }
    free(args);
    return NULL;
}

/*
 * Put a listener in a known state during
 * initialization.
 */
void
listener_init(struct ostp_listener *lp)
{
    lp->port = LISTEN_PORT;
    lp->on_recv = NULL;
}

int
listener_bind(struct ostp_listener *lp)
{
    struct ostp_session *session;
    struct sockaddr_in saddr;
    int error;

    lp->serv_sockfd = socket(AF_INET, SOCK_STREAM, 0);
    if (lp->serv_sockfd < 0) {
        perror("Failed to create socket\n");
        return -1;
    }

    saddr.sin_family = AF_INET;
    saddr.sin_addr.s_addr = INADDR_ANY;
    saddr.sin_port = htons(LISTEN_PORT);
    error = bind(lp->serv_sockfd, (struct sockaddr *)&saddr, sizeof(saddr));
    if (error < 0) {
        perror("Failed to bind socket\n");
        close(lp->serv_sockfd);
        return error;
    }

    if ((error = listen(lp->serv_sockfd, MAX_BACKLOG)) < 0) {
        perror("Failed to listen");
        return error;
    }

    return 0;
}

int
listener_poll(struct ostp_listener *lp)
{
    struct sockaddr_in caddr;
    struct recv_args *recv_ap;
    struct ostp_client *c;
    socklen_t caddr_len;
    pthread_t client_td;
    int client_sock, n, error = 0;
    unsigned char *session_key;
    char *ip, buf[4096];

    memset(lp->clients, -1, sizeof(lp->clients));
    lp->clients[0].sockfd = lp->serv_sockfd;
    lp->client_count = 1;

    while (1) {
        FD_ZERO(&lp->client_fds);

        for (int i = 0; i < MAX_CLIENTS; ++i) {
            if (lp->clients[i].sockfd >= 0)
                FD_SET(lp->clients[i].sockfd, &lp->client_fds);
        }

        if (select(1024, &lp->client_fds, NULL, NULL, NULL) < 0) {
            perror("select");
            continue;
        }

        /* Check if the servers socket has new connections */
        if (FD_ISSET(lp->serv_sockfd, &lp->client_fds)) {
            caddr_len = sizeof(caddr);
            client_sock = accept(lp->serv_sockfd, (struct sockaddr *)&caddr,
                &caddr_len);

            if (client_sock < 0) {
                perror("accept");
                continue;
            }

            if (lp->client_count >= MAX_CLIENTS) {
                printf("New connection rejected, max clients reached\n");
                close(client_sock);
                continue;
            }

            for (int i = 0; i < MAX_CLIENTS; ++i) {
                c = &lp->clients[i];
                if (c->sockfd < 0) {
                    c->sockfd = client_sock;
                    c->authed = 0;
                    ++lp->client_count;

                    ip = inet_ntoa(caddr.sin_addr);
                    memcpy(c->host, ip, strlen(ip) + 1);
                    printf("Incoming connection from %s\n", ip);
                    handle_client(&caddr, c, lp);
                    break;
                }
            }
        }

        /* Handle data from clients */
        for (int i = 1; i < MAX_CLIENTS; ++i) {
            c = &lp->clients[i];
            if (c->sockfd < 0 || !c->authed)
                continue;
            if (!FD_ISSET(c->sockfd, &lp->client_fds))
                continue;

            session_key = c->session.session_key;
            n = recv_frame(c->sockfd, sizeof(buf) - 1, session_key, buf);
            if (n < 0) {
                printf("recv_frame() failure, packet lost\n");
                continue;
            }
            if (n == 0) {
                printf("Peer disconnected\n");
                listener_close(lp, c);
                continue;
            }

            if (lp->on_recv != NULL) {
                recv_ap = malloc(sizeof(*recv_ap));
                if (recv_ap == NULL)
                    continue;

                /* Prepare on_recv() args */
                recv_ap->c = c;
                recv_ap->buf = buf;
                recv_ap->n = n;
                recv_ap->lp = lp;

                error = pthread_create(&recv_ap->td, NULL, recv_td, recv_ap);
                if (error != 0) {
                    free(recv_ap);
                    continue;
                }
            }
        }
    }

    close(client_sock);
    close(lp->serv_sockfd);
}

void
listener_cleanup(struct ostp_listener *lp)
{
    struct ostp_client *c;

    for (int i = 0; i < MAX_CLIENTS; ++i) {
        c = &lp->clients[i];
        listener_close(lp, c);
    }

    close(lp->serv_sockfd);
}

void
listener_close(struct ostp_listener *lp, struct ostp_client *c)
{
    if (lp->client_count == 0) {
        return;
    }

    close(c->sockfd);
    c->sockfd = -1;
    memset(&c->session, 0, sizeof(c->session));
    --lp->client_count;
}

/*
 * Get a list of clients from a listener.
 *
 * Returns the number of clients on success,
 * on failure this function returns a less than
 * zero value.
 *
 * @lp: Listener.
 * @res: An array of clients to be set.
 */
int
listener_clients(struct ostp_listener *lp, struct ostp_client **res)
{
    if (lp->client_count == 0 || lp->client_count > MAX_CLIENTS) {
        return -1;
    }

    for (int i = 0; i < lp->client_count + 1; ++i) {
        res[i] = &lp->clients[i];
    }

    return lp->client_count;
}