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

#include <string.h>
#include <unistd.h>
#include <stdio.h>
#include <error.h>
#include <stdint.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <pwd.h>
#include <getopt.h>

#define CAP_SLEEP(newval) do { if ( (newval) < sleepTime ) sleepTime = (newval); } while (0)

static time_t lastActivity;
static int userCount;

// Tracked user sessions
#define USERS (50)
static struct user *users;

// List of reboot and shutdown times
static const char *shutdownActions[SHUTDOWN_ENUM_END] = {
    [REBOOT] = "reboot",
    [POWEROFF] = "poweroff",
    [KEXEC] = "kexec",
    [SUSPEND] = "suspend",
};
#define TABMAX (20)
static struct {
    time_t deadline; // MONOTONIC; 0 = none, or last one was disarmed and time elapsed past deadline
    time_t noUserDeadline; // In case no user is logged in, consider this deadline too
    enum Shutdown action;
    bool disarmed;
    bool force;
} nextAction;

// Config via command line
static struct {
    char *shutdownCommand;
    struct time shutdown[TABMAX];
    int shutdowns;
    int poweroffTimeout;
    int logoutTimeout;
    int suspendTimeout;
    int saverTimeout;
    int dpmsTimeout;
    int gracePeriod;
    int minIdle;
    bool killUserProcesses;
    bool logoutIsShutdown;
} config;

static time_t combineTime( const time_t now, const struct time * time );

static void getXAuthority( struct user *user );

static bool parseCmdline( int argc, char **argv );

static void execShutdown( enum Shutdown action );

static void userLoggedOut(struct user* usr);

int main( int argc, char **argv )
{
    int idx;
    time_t startWall, startMono, endWall, endMono;
    init_time();
    if ( ! parseCmdline( argc, argv ) )
        return 2;
    // Always line buffered output
    setvbuf( stdout, NULL, _IOLBF, -BUFSIZ );
    setvbuf( stderr, NULL, _IOLBF, -BUFSIZ );
    if ( _testmode ) {
        fprintf( stderr, "Test mode only!\n" );
    }
    users = calloc( USERS, sizeof(*users) );
    nextAction.deadline = 0;
    startWall = time( NULL );
    startMono = now();
    // Calculate longest possible sleep time without delaying any timeout related actions
    // But cap at 60 seconds
    int defaultSleep = 60;
    if ( config.logoutTimeout > 0 && config.logoutTimeout < defaultSleep ) {
        defaultSleep = config.logoutTimeout;
    }
    if ( config.saverTimeout > 0 && config.saverTimeout < defaultSleep ) {
        defaultSleep = config.saverTimeout;
    }
    if ( config.dpmsTimeout > 0 && config.dpmsTimeout < defaultSleep ) {
        defaultSleep = config.dpmsTimeout;
    }
    int listenFd = rpc_open();
    if ( listenFd == -1 )
        return 1;
    lastActivity = now();
    for ( ;; ) {
        int sleepTime = defaultSleep;
        const int count = getUserList( users, USERS );
        userCount = count;
        for ( idx = 0; idx < count; ++idx ) {
            struct user * const usr = &users[idx];
            if ( !usr->online ) {
                if ( !usr->lastOnline ) {
                    fprintf( stderr, "This will never happen \\o/\n" );
                    continue;
                }
                // User was logged in, but isn't anymore
                userLoggedOut( usr );
                continue;
            }
            const time_t NOW = time( NULL );
            const time_t monoNOW = now();
            if ( *usr->display ) {
                // User has an X11 session going, get actual idle time
                struct x11state screen;
                getXAuthority( usr );
                if ( !getX11IdleTimes( usr->xauth, usr->display, &screen ) ) {
                    // Failed, assume active to be safe
                    printf( "X11 query failed. Using NOW.\n ");
                    usr->lastActivity = NOW;
                    usr->lastActivityOffset = 0;
                } else {
                    // Fiddle and fumble
                    usr->isLocked = ( screen.saverState == SAVER_LOCKED );
                    if ( screen.saverState == SAVER_OFF
                            || ( screen.saverState == SAVER_ACTIVE && usr->lockTime > 0 ) ) {
                        if ( usr->lockTime > 0 ) {
                            printf( "%s #### Session got unlocked (State: %d, LockTime: %u).\n",
                                    usr->display, (int)screen.saverState, (unsigned int)usr->lockTime );
                            // Session is not locked, unconditionally update timestamp
                            usr->lockTime = 0;
                            usr->lockTimeOffset = 0;
                        }
                        usr->lastActivity = NOW - screen.idleSeconds;
                        usr->lastActivityOffset = 0;
                    } else if ( usr->isLocked && screen.lockTimestamp != 0 && screen.lockTimestamp - usr->lockTime > 10 ) { // Allow for slight offset in case we manually set this
                        // Session is locked
                        if ( usr->lastActivity == 0 || usr->lockTime != 0 ) {
                            // We have no activity log yet, or quickly unlocked and relocked -- update idle time
                            usr->lastActivity = screen.lockTimestamp;
                            usr->lastActivityOffset = 0;
                            printf( "%s #### Session got externally re-locked.\n", usr->display );
                        } else {
                            printf( "%s #### Session got externally locked.\n", usr->display );
                        }
                        usr->lockTime = screen.lockTimestamp;
                        usr->lockTimeOffset = 0;
                    }
                    // If this file exists, the user doesn't want screen saver or DPMS to trigger
                    const bool useSaver = ( config.saverTimeout > 0 || config.dpmsTimeout > 0 )
                        && ( usr->saverFile[0] == '-' || access( usr->saverFile, F_OK ) == -1 );
                    const int idleTime = NOW - ( usr->lastActivity + usr->lastActivityOffset );
                    if ( config.dpmsTimeout > 0 && screen.screenStandby != SCREEN_UNKNOWN ) {
                        int want = SCREEN_UNKNOWN;
                        if ( config.logoutTimeout > 0 && idleTime + 300 > config.logoutTimeout ) {
                            want = SCREEN_ON;
                        } else if ( ! nextAction.disarmed && nextAction.deadline != 0 && nextAction.deadline - monoNOW < 300 ) {
                            want = SCREEN_ON;
                        } else if ( useSaver && idleTime > config.dpmsTimeout && screen.idleSeconds >= 60 ) {
                            want = SCREEN_OFF;
                        }
                        if ( want != SCREEN_UNKNOWN && screen.screenStandby != want ) {
                            printf( "#### Powering %s %s.\n", usr->display, want == SCREEN_ON ? "ON" : "OFF" );
                            setScreenDpms( usr->xauth, usr->display, want );
                        }
                    }
                    if ( useSaver && config.saverTimeout > 0 ) {
                        int want = SAVER_OFF;
                        if ( config.gracePeriod >= 0 && idleTime > config.saverTimeout + config.gracePeriod
                                && ( screen.saverState == SAVER_OFF || screen.saverState == SAVER_ACTIVE ) ) {
                            want = SAVER_LOCKED;
                        } else if ( idleTime > config.saverTimeout && screen.saverState == SAVER_OFF ) {
                            want = SAVER_ACTIVE;
                        } else if ( screen.saverState == SAVER_OFF ) {
                            CAP_SLEEP( config.saverTimeout - idleTime );
                        }
                        if ( want != SAVER_OFF ) {
                            enableScreenSaver( usr->xauth, usr->display, want );
                            usr->lockTime = NOW;
                            usr->lockTimeOffset = 0;
                            if ( screen.screenStandby == SCREEN_OFF ) {
                                // Screen was off before, but our ugly vmware ungrab hack might
                                // have simulated input and woke up the screen - enable standby
                                // again right away
                                setScreenDpms( usr->xauth, usr->display, SCREEN_OFF );
                            }
                        }
                    }
                } // End X11 handling
            }
            //printf( "Have User %d: %s Locked %d (%+d) Idle %d (%+d) -- Leader: %d\n", idx, usr->user, (int)usr->lockTime, usr->lockTimeOffset, (int)usr->lastActivity, usr->lastActivityOffset, (int)usr->sessionLeader );
            const int idleTime = NOW - usr->lastActivity;
            // See if we need to shorten sleep time for a pending session kill
            // or warn the user about an imminent session kill
            // or whether the screen saver is supposed to activate
            if ( config.logoutTimeout > 0 ) {
                // Idle logout is enabled
                const int remaining = config.logoutTimeout - idleTime;
                if ( remaining <= 2 ) {
                    // Kill session, or if enabled, shut down machine entirely, if this was the main session
                    if ( config.logoutIsShutdown && usr->display[0] == ':' && usr->display[1] == '0' ) {
                        execShutdown( POWEROFF );
                    } else {
                        killSession( usr );
                    }
                } else if ( remaining <= 65 ) {
                    warn_showDefaultWarning( usr, WARN_LOGOUT, remaining );
                    CAP_SLEEP( remaining );
                } else if ( remaining < 310 ) {
                    if ( remaining % 30 < 10 ) {
                        warn_showDefaultWarning( usr, WARN_LOGOUT_LOW, remaining );
                    }
                    CAP_SLEEP( ( remaining % 30 ) - 2 );
                } else {
                    CAP_SLEEP( remaining - 305 );
                }
                usr->logoutTime = NOW + remaining;
            }
            // Update timestamp of last use of system
            const time_t laMono = monoNOW - idleTime;
            if ( laMono > lastActivity ) {
                lastActivity = laMono;
            }
        } // End loop over users
        const time_t monoNOW = now();
        // See if any reboot or poweroff is due (if none is currently pending or still 5min+ away)
        if ( ! nextAction.force && config.shutdowns > 0 // Next action is not forced, and have shutdown schedule
                // and no current action pending, or next action more than 5 mins away
                && ( nextAction.deadline == 0 || nextAction.deadline - monoNOW > 300 ) ) {
            // next action is more than 5 mins away, reset for now
            if ( nextAction.deadline != 0 && nextAction.deadline - monoNOW > 300 ) {
                nextAction.deadline = 0;
            }
            const time_t NOW = time( NULL );
            time_t deadline;
            int delta;
            // Loop through all scheduled actions and find the next one
            for ( idx = 0; idx < config.shutdowns; ++idx ) {
                deadline = combineTime( NOW, &config.shutdown[idx] );
                if ( deadline == 0 )
                    continue;
                delta = deadline - NOW;
                if ( delta >= -6 && ( nextAction.deadline == 0 || ( nextAction.deadline - monoNOW - 3 ) > delta ) ) {
                    // Engage
                    nextAction.deadline = monoNOW + ( delta < 0 ? 0 : delta );
                    nextAction.noUserDeadline = nextAction.deadline;
                    nextAction.disarmed = false;
                    nextAction.action = config.shutdown[idx].action;
                    if ( delta < 300 ) {
                        printf( "Scheduling reboot/poweroff in %d seconds\n", delta );
                    }
                    CAP_SLEEP( delta );
                }
            }
        }
        // Check unused idle time
        if ( count == 0 ) {
            if ( config.poweroffTimeout != 0 ) {
                int delta = config.poweroffTimeout - ( monoNOW - lastActivity );
                if ( delta <= 0 ) {
                    execShutdown( POWEROFF );
                } else {
                    CAP_SLEEP( delta );
                }
            }
            if ( config.suspendTimeout != 0 ) {
                int delta = config.suspendTimeout - ( monoNOW - lastActivity );
                if ( delta <= 0 ) {
                    execShutdown( SUSPEND );
                } else {
                    CAP_SLEEP( delta );
                }
            }
        }
        if ( nextAction.deadline != 0 && ( count == 0 || ( monoNOW - lastActivity ) >= config.minIdle || nextAction.force ) ) {
            // Some action seems pending
            int remaining;
            // No user logged in, no user deadline is max. 300 seconds in the past -> use that value
            if ( userCount == 0 && nextAction.noUserDeadline + 300 > monoNOW && nextAction.noUserDeadline < nextAction.deadline ) {
                remaining = nextAction.noUserDeadline - monoNOW;
                if ( remaining < 0 ) {
                    remaining = 0;
                }
            } else {
                remaining = nextAction.deadline - monoNOW;
            }
            if ( remaining < -60 ) {
                if ( ! nextAction.disarmed ) {
                    fprintf( stderr, "Missed planned action!? Late by %d seconds. Ignored.\n", -remaining );
                }
                nextAction.deadline = 0;
                nextAction.noUserDeadline = 0;
            } else if ( ! nextAction.disarmed ) {
                if ( remaining <= 1 ) {
                    // Execute!
                    execShutdown( nextAction.action );
                    nextAction.disarmed = true;
                    nextAction.force = false;
                } else if ( remaining < 310 ) {
                    if ( remaining % 30 < 8 ) {
                        enum Warning w = WARN_REBOOT;
                        if ( nextAction.action == POWEROFF ) {
                            w = WARN_POWEROFF;
                        }
                        for ( idx = 0; idx < count; ++idx ) {
                            warn_showDefaultWarning( &users[idx], w, remaining );
                        }
                    }
                    CAP_SLEEP( ( remaining % 15 ) - 2 );
                } else {
                    CAP_SLEEP( remaining - 305 );
                }
            }
        }
        do {
            // Handle requests
            const time_t oldDeadline = nextAction.deadline;
            // Sleep until next run
            int sn = sleepTime - ( endMono - startMono );
            //printf( "Sleeping %d seconds\n ", sn );
            if ( rpc_wait( listenFd, sn < 1 ? 1 : sn ) ) {
                rpc_handle( listenFd );
                if ( nextAction.deadline != oldDeadline ) {
                    // Might have set a new scheduled action, run main loop again soon
                    CAP_SLEEP( 1 );
                } else {
                    // Otherwise, still update rather quickly, in case we get polled again
                    CAP_SLEEP( 5 );
                }
            }
            // Detect time changes
            endWall = time( NULL );
            endMono = now();
            // Sloppy way of preventing we loop too fast -- handle rpc callbacks but don't run main logic
        } while ( sleepTime > ( endMono - startMono ) );
        // Adjust for clock changes
        const int diff = ( endWall - startWall ) - ( endMono - startMono );
        if ( diff < -2 || diff > 2 ) {
            printf( "Correcting time by %d seconds\n", diff );
            for ( idx = 0; idx < count; ++idx ) {
                if ( users[idx].display[0] == '\0' )
                    continue;
                if ( users[idx].lockTime != 0 ) {
                    users[idx].lockTimeOffset += diff;
                }
                users[idx].lastActivityOffset += diff;
            }
        }
        startWall = endWall;
        startMono = endMono;
    }
    return 0;
}

static time_t combineTime( const time_t now, const struct time * time )
{
    struct tm dtime;
    if ( localtime_r( &now, &dtime ) == NULL )
        return 0;
    time_t result;
    dtime.tm_hour = time->hour;
    dtime.tm_min = time->minute;
    dtime.tm_sec = 0;
    result = mktime( &dtime );
    if ( result + 6 < now ) {
        // Try again tomorrow
        dtime.tm_mday++;
        result = mktime( &dtime );
    }
    if ( result == -1 ) {
        perror( "Cannot convert struct tm to time_t via mktime" );
        return 0;
    }
    if ( result < now ) {
        fprintf( stderr, "combineTime: Even +1 day seems in the past!?\n" );
        return 0;
    }
    return result;
}

/**
 * This doesn't really use any clever logic but rather assumes that
 * it will be $HOME/.Xauthority
 */
static void getXAuthority( struct user *user )
{
    if ( user->xauth[0] != '\0' )
        return;
    char buffer[1000];
    struct passwd p, *ok;
    if ( getpwnam_r( user->user, &p, buffer, sizeof(buffer), &ok ) != 0 || ok == NULL ) {
        user->xauth[0] = '-';
        user->xauth[1] = '\0';
        user->saverFile[0] = '-';
        user->saverFile[1] = '\0';
        return;
    }
    // Got valid data
    snprintf( user->saverFile, SAVERLEN, "%s/.no-saver", p.pw_dir );
    snprintf( user->xauth, AUTHLEN, "%s/.Xauthority", p.pw_dir );
    if ( access( user->xauth, F_OK ) == -1 ) {
        user->xauth[0] = '-';
        user->xauth[1] = '\0';
    }
}

// cmdline options. Only support long for now.
static struct option long_options[] = {
    { "reboot", required_argument, NULL, 'r' },
    { "poweroff", required_argument, NULL, 'p' },
    { "poweroff-timeout", required_argument, NULL, 'pot' },
    { "suspend-timeout", required_argument, NULL, 'sbt' },
    { "logout-timeout", required_argument, NULL, 'lot' },
    { "screensaver-timeout", required_argument, NULL, 'sst' },
    { "dpms-timeout", required_argument, NULL, 'dpms' },
    { "grace-period", required_argument, NULL, 'gp' },
    { "min-idle", required_argument, NULL, 'min' },
    { "cmd", required_argument, NULL, 'cmd' },
    { "send", required_argument, NULL, 'send' },
    { "kill-user-processes", no_argument, NULL, 'kill' },
    { "logout-is-shutdown", no_argument, NULL, 'lis' },
    { "test", no_argument, NULL, 't' },
    { NULL, 0, NULL, 0 }
};

static bool parseTime( const char *str, struct time *result )
{
    char *next;
    result->hour = (int)strtol( str, &next, 10 );
    if ( result->hour < 0 || result->hour > 24 )
        return false;
    result->hour %= 24; // Allow 24:00
    if ( *next != ':' )
        return false;
    result->minute = (int)strtol( next + 1, &next, 10 );
    if ( result->minute < 0 || result->minute > 59 )
        return false;
    if ( *next != '\0' ) {
        fprintf( stderr, "Ignoring trailing garbage after minute\n" );
    }
    return true;
}

/**
 * Parse command line and fill all the tables / vars etc.
 * Return false if command line is unparsable.
 */
static bool parseCmdline( int argc, char **argv )
{
    int ch;
    config.shutdowns = 0;
    config.logoutTimeout = 0;
    config.poweroffTimeout = 0;
    config.suspendTimeout = 0;
    config.saverTimeout = 0;
    config.dpmsTimeout = 0;
    config.gracePeriod = -1;
    config.shutdownCommand = NULL;
    config.minIdle = 0;
    config.killUserProcesses = false;
    config.logoutIsShutdown = false;
    while ( ( ch = getopt_long( argc, argv, "", long_options, NULL ) ) != -1 ) {
        switch ( ch ) {
        case 'pot':
            config.poweroffTimeout = atoi( optarg );
            break;
        case 'sbt':
            config.suspendTimeout = atoi( optarg );
            break;
        case 'lot':
            config.logoutTimeout = atoi( optarg );
            break;
        case 'sst':
            config.saverTimeout = atoi( optarg );
            break;
        case 'dpms':
            config.dpmsTimeout = atoi( optarg );
            break;
        case 'gp':
            config.gracePeriod = atoi( optarg );
            break;
        case 'min':
            config.minIdle = atoi( optarg );
            break;
        case 'r':
        case 'p':
            if ( config.shutdowns < TABMAX ) {
                if ( parseTime( optarg, &config.shutdown[config.shutdowns] ) ) {
                    config.shutdown[config.shutdowns].action = ( ch == 'r' ? REBOOT : POWEROFF );
                    config.shutdowns++;
                } else {
                    fprintf( stderr, "Could not parse shutdown time %s\n", optarg );
                }
            } else {
                fprintf( stderr, "Ignoring shutdown time %s: Table full\n", optarg );
            }
            break;
        case 't':
            _testmode = 1;
            break;
        case 'cmd':
            config.shutdownCommand = strdup( optarg );
            break;
        case 'send':
            exit( !rpc_send( optarg ) );
            break;
        case 'kill':
            config.killUserProcesses = true;
            break;
        case 'lis':
            config.logoutIsShutdown = true;
            break;
        default:
            fprintf( stderr, "Unhandled command line option %d, aborting\n", ch );
            return false;
        }
    }
    return true;
}

static void execShutdown( enum Shutdown action )
{
    if ( action < 0 && action >= SHUTDOWN_ENUM_END ) {
        fprintf( stderr, "Invalid shutdown action %d\n", (int)action );
        return;
    }
    if ( _testmode ) {
        printf( "[dryrun] Not execution shutdown(%d)\n", (int)action );
        return;
    }
    if ( config.shutdownCommand != NULL ) {
        printf( "Executing shutdown via %s %s\n", config.shutdownCommand, shutdownActions[action] );
        run( true, config.shutdownCommand, shutdownActions[action], (char*)NULL );
        return;
    }
    // Builtin
    printf( "Executing shutdown via builtin handler for %s\n", shutdownActions[action] );
    switch ( action ) {
    case REBOOT:
    case POWEROFF:
    case SUSPEND:
        run( true, "systemctl", shutdownActions[action], (char*)NULL );
        break;
    case KEXEC:
        run( true, "systemctl", "isolate", "kexec.target", (char*)NULL );
        break;
    default:
        fprintf( stderr, "BUG! Unhandled shutdown action %d\n", (int)action );
        break;
    }
}

void main_getStatus( const char **nextString, time_t *deadline, int *numUsers )
{
    if ( nextAction.deadline == 0 || nextAction.disarmed ) {
        *deadline = 0;
        return;
    }
    *deadline = ( nextAction.deadline - now() ) + time( NULL );
    *nextString = shutdownActions[nextAction.action];
    *numUsers = userCount;
}

struct user* main_getUser( const char *terminal )
{
    for ( int i = 0; i < userCount; ++i ) {
        if ( strcmp( users[i].device, terminal ) == 0
                || strcmp( users[i].display, terminal ) == 0 ) {
            return &users[i];
        }
    }
    return NULL;
}

/**
 * Called from RPC, only by root: Set the next action (but only
 * if it's sooner than anything already pending)
 */
void main_queueAction( enum Shutdown action, int delta )
{
    if ( delta < 0 || action < 0 || action >= SHUTDOWN_ENUM_END )
        return;
    time_t monoNOW = now();
    int realDelta = delta;
    if ( userCount != 0 ) {
        int idleTime = monoNOW - lastActivity;
        if ( idleTime < 3600 ) {
            if ( delta < 306 ) {
                delta = 306;
            }
        } else if ( idleTime < 7200 ) {
            if ( delta < 66 ) {
                delta = 66;
            }
        }
    }
    if ( nextAction.deadline == 0 || ( nextAction.deadline - monoNOW - 3 ) > delta ) {
        // Engage
        nextAction.deadline = monoNOW + delta;
        nextAction.noUserDeadline = monoNOW + realDelta;
        nextAction.disarmed = false;
        nextAction.force = true;
        nextAction.action = action;
        printf( "RPC: Scheduling reboot/poweroff in %d seconds\n", delta );
    }
}

void main_warnAll( const char *message )
{
    for ( int idx = 0; idx < userCount; ++idx ) {
        warn_showCustomWarning( &users[idx], "Warning", message );
    }
}

static void userLoggedOut(struct user* usr)
{
    if ( !config.killUserProcesses )
        return;
    for ( int i = 0; i < userCount; ++i ) {
        if (users[i].online && strcmp(users[i].user, usr->user) == 0)
            return; // Still an active session
    }
    struct passwd *u = getpwnam( usr->user );
    if ( u == NULL || u->pw_uid < 1000 )
        return; // Ignore system users
    printf( "Killing remaining processes of %s\n", usr->user );
    run( true, "pkill", "-u", usr->user, (char*)NULL );
}