/*
Written by retoor@molodetz.nl

Tikker - Keystroke monitoring and productivity analytics application.
Captures keyboard input events, stores in SQLite, provides statistics.

MIT License

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
*/

#include "sormc.h"
#include "tikker_types.h"
#include "tikker_db.h"
#include "tikker_decode.h"
#include "tikker_words.h"
#include "tikker_stats.h"

#include <fcntl.h>
#include <linux/input.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <unistd.h>

static int tikker_db = 0;

static void tikker_init_pragmas(void) {
    sormq(tikker_db, "PRAGMA journal_mode=WAL");
    sormq(tikker_db, "PRAGMA synchronous=NORMAL");
    sormq(tikker_db, "PRAGMA cache_size=-65536");
    sormq(tikker_db, "PRAGMA temp_store=MEMORY");
    sormq(tikker_db, "PRAGMA mmap_size=268435456");
}

static void tikker_init_schema(void) {
    sormq(tikker_db, "CREATE TABLE IF NOT EXISTS kevent ("
        "id INTEGER PRIMARY KEY AUTOINCREMENT, "
        "code, event, name, timestamp, char)");
    sormq(tikker_db, "CREATE INDEX IF NOT EXISTS idx_kevent_event ON kevent(event)");
    sormq(tikker_db, "CREATE INDEX IF NOT EXISTS idx_kevent_timestamp ON kevent(timestamp)");
    sormq(tikker_db, "CREATE INDEX IF NOT EXISTS idx_kevent_code ON kevent(code)");
    sormq(tikker_db, "CREATE INDEX IF NOT EXISTS idx_kevent_event_code ON kevent(event, code)");
    sormq(tikker_db, "CREATE INDEX IF NOT EXISTS idx_kevent_event_timestamp ON kevent(event, timestamp)");
    sormq(tikker_db, "CREATE INDEX IF NOT EXISTS idx_kevent_event_timestamp_char ON kevent(event, timestamp, char)");

    sormq(tikker_db, "CREATE TABLE IF NOT EXISTS kevent_daily ("
        "date TEXT PRIMARY KEY, "
        "pressed INTEGER DEFAULT 0, "
        "released INTEGER DEFAULT 0, "
        "repeated INTEGER DEFAULT 0)");

    sormq(tikker_db, "CREATE TABLE IF NOT EXISTS kevent_hourly ("
        "date_hour TEXT PRIMARY KEY, "
        "pressed INTEGER DEFAULT 0)");

    sormq(tikker_db, "CREATE TABLE IF NOT EXISTS kevent_key_counts ("
        "code INTEGER PRIMARY KEY, "
        "count INTEGER DEFAULT 0)");
}

static char *tikker_resolve_device_name(int fd) {
    static char device_name[256];
    device_name[0] = 0;
    if (ioctl(fd, EVIOCGNAME(sizeof(device_name)), device_name) < 0) {
        return NULL;
    }
    return device_name;
}

static void tikker_populate_cache(void) {
    sorm_ptr count = sormq(tikker_db, "SELECT COUNT(*) FROM kevent_daily");
    if (count) {
        char *csv = (char *)count;
        char *line = strchr(csv, '\n');
        if (line) line++;
        else line = csv;
        int cached_days = atoi(line);
        free(count);
        if (cached_days > 0) return;
    }

    sormq(tikker_db, "INSERT OR REPLACE INTO kevent_daily (date, pressed, released, repeated) "
        "SELECT STRFTIME('%%Y-%%m-%%d', timestamp) as date, "
        "SUM(CASE WHEN event = 'PRESSED' THEN 1 ELSE 0 END), "
        "SUM(CASE WHEN event = 'RELEASED' THEN 1 ELSE 0 END), "
        "SUM(CASE WHEN event = 'REPEATED' THEN 1 ELSE 0 END) "
        "FROM kevent GROUP BY date");

    sormq(tikker_db, "INSERT OR REPLACE INTO kevent_hourly (date_hour, pressed) "
        "SELECT STRFTIME('%%Y-%%m-%%d.%%H', timestamp), COUNT(*) "
        "FROM kevent WHERE event = 'PRESSED' GROUP BY 1");

    sormq(tikker_db, "INSERT OR REPLACE INTO kevent_key_counts (code, count) "
        "SELECT code, COUNT(*) FROM kevent WHERE event = 'PRESSED' GROUP BY code");
}

static void tikker_update_cache_for_event(const char *event, int code) {
    sormq(tikker_db, "INSERT INTO kevent_daily (date, pressed, released, repeated) "
        "VALUES (DATE('now'), 0, 0, 0) "
        "ON CONFLICT(date) DO NOTHING");

    if (strcmp(event, "PRESSED") == 0) {
        sormq(tikker_db, "UPDATE kevent_daily SET pressed = pressed + 1 WHERE date = DATE('now')");

        sormq(tikker_db, "INSERT INTO kevent_hourly (date_hour, pressed) "
            "VALUES (STRFTIME('%%Y-%%m-%%d.%%H', 'now'), 1) "
            "ON CONFLICT(date_hour) DO UPDATE SET pressed = pressed + 1");

        sormq(tikker_db, "INSERT INTO kevent_key_counts (code, count) VALUES (%d, 1) "
            "ON CONFLICT(code) DO UPDATE SET count = count + 1", code);
    } else if (strcmp(event, "RELEASED") == 0) {
        sormq(tikker_db, "UPDATE kevent_daily SET released = released + 1 WHERE date = DATE('now')");
    } else {
        sormq(tikker_db, "UPDATE kevent_daily SET repeated = repeated + 1 WHERE date = DATE('now')");
    }
}

static void tikker_print_usage(void) {
    printf("Usage: tikker [command] [options]\n\n");
    printf("Commands:\n");
    printf("  (no command)              Start keyboard monitoring\n");
    printf("  presses_today             Show today's keystroke count\n");
    printf("  stats daily               Daily keystroke statistics\n");
    printf("  stats hourly [DATE]       Hourly breakdown (default: today)\n");
    printf("  stats weekly              Weekly keystroke statistics\n");
    printf("  stats weekday             Weekday comparison\n");
    printf("  stats top-keys [N]        Top N keys (default: 10)\n");
    printf("  stats top-words [N]       Top N words (default: 10)\n");
    printf("  stats summary             Overall summary statistics\n");
    printf("  export                    Export logs to logs_plain/\n");
    printf("  decode [FILE]             Decode keystroke log file\n");
    printf("\nOptions:\n");
    printf("  --device='NAME'           Monitor specific device\n");
}

static int tikker_handle_stats_command(int argc, char *argv[]) {
    if (argc < 1) {
        tikker_print_usage();
        return 1;
    }

    const char *cmd = argv[0];

    if (strcmp(cmd, "daily") == 0) {
        return tikker_stats_daily(tikker_db);
    }
    if (strcmp(cmd, "hourly") == 0) {
        const char *date = (argc > 1) ? argv[1] : tikker_get_today_date();
        return tikker_stats_hourly(tikker_db, date);
    }
    if (strcmp(cmd, "weekly") == 0) {
        return tikker_stats_weekly(tikker_db);
    }
    if (strcmp(cmd, "weekday") == 0) {
        return tikker_stats_weekday(tikker_db);
    }
    if (strcmp(cmd, "top-keys") == 0) {
        int limit = (argc > 1) ? atoi(argv[1]) : 10;
        if (limit <= 0) limit = 10;
        return tikker_stats_top_keys(tikker_db, limit);
    }
    if (strcmp(cmd, "top-words") == 0) {
        int limit = (argc > 1) ? atoi(argv[1]) : 10;
        if (limit <= 0) limit = 10;
        return tikker_stats_top_words(tikker_db, limit);
    }
    if (strcmp(cmd, "summary") == 0) {
        return tikker_stats_summary(tikker_db);
    }

    fprintf(stderr, "Unknown stats command: %s\n", cmd);
    tikker_print_usage();
    return 1;
}

static int tikker_run_keylogger(int argc, char *argv[]) {
    char *device_to_read = rargs_get_option_string(argc, argv, "--device", TIKKER_DEVICE_DEFAULT);

    ulonglong times_repeated = 0;
    ulonglong times_pressed = 0;
    ulonglong times_released = 0;

    int keyboard_fds[TIKKER_MAX_DEVICES];
    int num_keyboards = 0;

    for (int i = 0; i < TIKKER_MAX_DEVICES; i++) {
        char device_path[32];
        snprintf(device_path, sizeof(device_path), "%s%d", TIKKER_DEVICE_PATH, i);
        int fd = open(device_path, O_RDONLY);
        if (fd < 0) continue;

        char *device_name = tikker_resolve_device_name(fd);
        if (!device_name) {
            close(fd);
            continue;
        }

        bool is_target = strstr(device_name, device_to_read) != NULL;
        printf("[%s] %s. Mount: %s.\n", is_target ? "-" : "+", device_name, device_path);

        if (is_target) {
            keyboard_fds[num_keyboards++] = fd;
        } else {
            close(fd);
        }
    }

    if (num_keyboards == 0) {
        fprintf(stderr, "No keyboard found. Are you running as root?\n"
                        "If your device is listed above with a minus [-] in front,\n"
                        "run this application using --device='[DEVICE_NAME]'\n");
        return 1;
    }

    printf("Monitoring %d keyboards.\n", num_keyboards);

    struct input_event ev;
    fd_set read_fds;

    while (1) {
        FD_ZERO(&read_fds);
        int max_fd = -1;
        for (int i = 0; i < num_keyboards; i++) {
            FD_SET(keyboard_fds[i], &read_fds);
            if (keyboard_fds[i] > max_fd) max_fd = keyboard_fds[i];
        }

        if (select(max_fd + 1, &read_fds, NULL, NULL, NULL) < 0) {
            perror("select error");
            break;
        }

        for (int i = 0; i < num_keyboards; i++) {
            if (!FD_ISSET(keyboard_fds[i], &read_fds)) continue;

            ssize_t bytes = read(keyboard_fds[i], &ev, sizeof(struct input_event));
            if (bytes != sizeof(struct input_event)) continue;
            if (ev.type != EV_KEY) continue;

            char *char_name = NULL;
            if (ev.code < sizeof(tikker_keycode_to_char) / sizeof(tikker_keycode_to_char[0])) {
                char_name = (char *)tikker_keycode_to_char[ev.code];
            }

            char keyboard_name[256];
            ioctl(keyboard_fds[i], EVIOCGNAME(sizeof(keyboard_name)), keyboard_name);

            char *event_name;
            if (ev.value == 1) {
                event_name = "PRESSED";
                times_pressed++;
            } else if (ev.value == 0) {
                event_name = "RELEASED";
                times_released++;
            } else {
                event_name = "REPEATED";
                times_repeated++;
            }

            sormq(tikker_db, "INSERT INTO kevent (code, event, name, timestamp, char) "
                "VALUES (%d, %s, %s, DATETIME('now'), %s)",
                ev.code, event_name, keyboard_name, char_name);

            tikker_update_cache_for_event(event_name, ev.code);

            printf("Keyboard: %s, Event: %s, Key Code: %d, Name: %s, "
                   "Pr: %lld Rel: %lld Rep: %lld\n",
                   keyboard_name, event_name, ev.code, char_name,
                   times_pressed, times_released, times_repeated);
        }
    }

    for (int i = 0; i < num_keyboards; i++) {
        close(keyboard_fds[i]);
    }

    return 0;
}

int main(int argc, char *argv[]) {
    tikker_db = sormc(TIKKER_DATABASE_NAME);
    tikker_init_pragmas();
    tikker_init_schema();
    tikker_populate_cache();

    if (argc >= 2) {
        if (strcmp(argv[1], "presses_today") == 0) {
            return tikker_stats_presses_today(tikker_db);
        }

        if (strcmp(argv[1], "stats") == 0) {
            if (argc < 3) {
                tikker_print_usage();
                return 1;
            }
            return tikker_handle_stats_command(argc - 2, argv + 2);
        }

        if (strcmp(argv[1], "export") == 0) {
            return tikker_export_logs(tikker_db);
        }

        if (strcmp(argv[1], "decode") == 0) {
            if (argc < 3) {
                fprintf(stderr, "Error: decode requires a filename\n");
                tikker_print_usage();
                return 1;
            }
            return tikker_decode_file(argv[2]);
        }

        if (strcmp(argv[1], "help") == 0 ||
            strcmp(argv[1], "--help") == 0 ||
            strcmp(argv[1], "-h") == 0) {
            tikker_print_usage();
            return 0;
        }

        if (strncmp(argv[1], "--device=", 9) == 0) {
            return tikker_run_keylogger(argc, argv);
        }

        fprintf(stderr, "Unknown command: %s\n", argv[1]);
        tikker_print_usage();
        return 1;
    }

    return tikker_run_keylogger(argc, argv);
}