#include "rpc.h"
#include "util.h"
#include "main.h"
#include "userlist.h"

#include <poll.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/un.h>
#include <stdio.h>
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdlib.h>

#define SENDSTRING(fd,x) write( fd, x, strlen(x) )

#define SOCKPATH "/run/idle-daemon"

static union {
    struct __attribute__((packed)) {
        int read;
        int write;
    };
    int array[2];
} rpcPipe = {
    .read = -1,
    .write = -1,
};

_Static_assert( sizeof(rpcPipe) == ( sizeof(int) * 2 ), "Structsize mismatch" );

static void handleClient( int fd, struct ucred *user );

int rpc_open( void )
{
    int fd = socket( AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC | SOCK_NONBLOCK, 0 );
    if ( fd == -1 ) {
        perror( "Cannot create local RPC socket" );
        return -1;
    }
    struct sockaddr_un address = {
        .sun_family = AF_UNIX,
        .sun_path = SOCKPATH,
    };
    unlink( SOCKPATH );
    if ( bind( fd, (struct sockaddr *)&address, sizeof(address) ) == -1 ) {
        perror( "Could not bind RPC socket" );
        close( fd );
        return -1;
    }
    chmod( SOCKPATH, 0777 );
    if ( listen( fd, 10 ) == -1 ) {
        perror( "Could not listen() on RPC socket" );
        close( fd );
        return -1;
    }
    if ( rpcPipe.write == -1 && rpcPipe.read == -1 ) {
        if ( pipe2( rpcPipe.array, O_DIRECT ) == -1 ) {
            perror( "Cannot create local RPC pipe" );
        } else {
            // Read end nonblocking
            int flags = fcntl( rpcPipe.read, F_GETFL, 0 );
            fcntl( rpcPipe.read, F_SETFL, flags | O_NONBLOCK );
        }
    }
    return fd;
}

bool rpc_wait( int listenFd, int seconds )
{
    struct pollfd pd[2] = {
        { .fd = listenFd, .events = POLLIN | POLLHUP | POLLRDHUP },
        { .fd = rpcPipe.read, .events = POLLIN | POLLHUP | POLLRDHUP },
    };
    int ret = poll( pd, 2, seconds * 1000 );
    if ( ret == -1 ) {
        perror( "Error polling RPC sockets" );
    }
    return ret > 0;
}

void rpc_handle( int listenFd )
{
    int fd;
    while ( ( fd = accept( listenFd, NULL, NULL ) ) != -1 ) {
        // Determine who connected
        socklen_t len;
        struct ucred ucred;
        len = sizeof(struct ucred);
        if ( getsockopt( fd, SOL_SOCKET, SO_PEERCRED, &ucred, &len ) == -1 ) {
            perror( "Could not get credentials of connection" );
            close( fd ); // TODO: Allow more for root (eg. cancel reboot/poweroff)
            continue;
        }
        if ( ! doublefork() ) {
            close( fd );
            continue; // Parent continues
        }
        handleClient( fd, &ucred );
        exit( 0 );
    }
    if ( errno != EAGAIN && errno != EWOULDBLOCK ) {
        perror( "accept() on RPC socket failed" );
    }
    // Check if any RPC client talks to us
    char buffer[200];
    ssize_t len;
    while ( ( len = read( rpcPipe.read, buffer, sizeof(buffer) - 1 ) ) > 0 ) {
        buffer[len] = '\0';
        const char *ptr = NULL;
        enum Shutdown action = SHUTDOWN_ENUM_END;
        if ( strncmp( buffer, "reboot ", 7 ) == 0 ) {
            ptr = buffer + 7;
            action = REBOOT;
        } else if ( strncmp( buffer, "poweroff ", 9 ) == 0 ) {
            ptr = buffer + 9;
            action = POWEROFF;
        } else if ( strncmp( buffer, "kexec ", 6 ) == 0 ) {
            ptr = buffer + 6;
            action = KEXEC;
        } else {
            fprintf( stderr, "Unknown RPC pipe callback: %s\n", buffer );
        }
        if ( ptr != NULL ) {
            char *end;
            int seconds = strtol( ptr, &end, 10 );
            if ( end != ptr && seconds >= 0 ) {
                main_queueAction( action, seconds );
            } else {
                fprintf( stderr, "RPC: Ignoring reboot/poweroff with invalid timeout\n" );
            }
        }
    }
    if ( len == -1 && errno != EAGAIN && errno != EWOULDBLOCK ) {
        perror( "Error reading vom RPC local pipe" );
    }
}

static void handleClient( int fd, struct ucred *creds )
{
    //printf( "Credentials from SO_PEERCRED: pid=%ld, euid=%ld, egid=%ld\n",
    //    (long) ucred.pid, (long) ucred.uid, (long) ucred.gid );
    // Make socket blocking (should be default on linux after accept() but...)
    int flags = fcntl( fd, F_GETFL, 0 );
    if ( flags != -1 ) {
        fcntl( fd, F_SETFL, flags & ~O_NONBLOCK );
    }
    // But set timeouts
    struct timeval tv = {
        .tv_sec = 2,
    };
    setsockopt( fd, SOL_SOCKET, SO_RCVTIMEO, (const char*)&tv, sizeof(tv) );
    setsockopt( fd, SOL_SOCKET, SO_SNDTIMEO, (const char*)&tv, sizeof(tv) );
    // Now read request
    char buffer[1000];
    ssize_t len = read( fd, buffer, sizeof(buffer) - 1 );
    if ( len <= 0 )
        return;
    buffer[len] = '\0';
    if ( strncmp( buffer, "get", 3 ) == 0 ) {
        // Get request
        FILE *s = fdopen( fd, "w" );
        if ( s == NULL ) {
            perror( "Cannot wrap socket in stream" );
            return;
        }
        // 1) Global state
        time_t deadline;
        const char *name = "none";
        int userCount;
        main_getStatus( &name, &deadline, &userCount );
        fprintf( s, "[General]\n"
                "nextAction=%s\n"
                "nextActionTime=%lld\n"
                "userCount=%d\n",
                name, (long long)deadline, userCount );
        // 2) Requested sessions
        char *tok = strtok( buffer + 4, " \t\n\r" );
        while ( tok != NULL ) {
            struct user *user = main_getUser( tok );
            if ( user != NULL ) {
                const int idleTime = (int)( time( NULL ) - ( user->lastActivity + user->lastActivityOffset ) );
                fprintf( s, "[%s]\n"
                        "logoutTime=%lld\n"
                        "locked=%d\n"
                        "idleSeconds=%d\n",
                        tok, (long long)user->logoutTime, (int)user->isLocked, idleTime );
            }
            tok = strtok( NULL, " \t\n\r" );
        }
        fflush( s );
        // Don't fclose since we need to shutdown below. This is a leak, but since we're
        // running in a child that will exit soon, just close your eyes for a second...
    } else if ( strncmp( buffer, "reboot ", 7 ) == 0 || strncmp( buffer, "poweroff ", 9 ) == 0
            || strncmp( buffer, "kexec ", 6 ) == 0 ) {
        if ( creds->uid != 0 ) {
            SENDSTRING( fd, "error only root can do this" );
        } else {
            if ( write( rpcPipe.write, buffer, len ) == -1 ) {
                perror( "RPC Child: Cannot write to parent pipe" );
                SENDSTRING( fd, "error cannot write to parent" );
            } else {
                SENDSTRING( fd, "ok" );
            }
        }
    } else if ( strncmp( buffer, "warn ", 5 ) == 0 ) {
        if ( creds->uid != 0 ) {
            SENDSTRING( fd, "error only root can do this" );
        } else {
            main_warnAll( buffer + 5 );
            SENDSTRING( fd, "ok" );
        }
    } else {
        SENDSTRING( fd, "error unknown command" );
    }
    shutdown( fd, SHUT_WR );
    read( fd, buffer, sizeof(buffer) );
    close( fd );
}

bool rpc_send( const char *data )
{
    struct sockaddr_un addr = {
        .sun_family = AF_UNIX,
        .sun_path = SOCKPATH,
    };
    int fd = socket( AF_UNIX, SOCK_STREAM, 0 );
    if ( fd == -1 ) {
        perror( "RPC Client: socket() failed" );
        return false;
    }
    if ( connect( fd, &addr, sizeof(addr) ) == -1 ) {
        perror( "RPC Client: connect() failed" );
        close( fd );
        return false;
    }
    if ( write( fd, data, strlen( data ) ) <= 0 ) {
        perror( "RPC Client: write() failed" );
        close( fd );
        return false;
    }
    shutdown( fd, SHUT_WR );
    char buffer[200];
    ssize_t len;
    if ( ( len = read( fd, buffer, sizeof(buffer) - 1 ) ) == -1 ) {
        perror( "RPC Client: read()ing reply failed" );
    } else {
        buffer[len] = '\0';
        printf( "%s\n", buffer );
    }
    close( fd );
    return true;
}