/*
 * vault.c -- Crypto Security Token Vault Gateway
 *
 * SPDX-License-Identifier: GPL-2.0-or-later
 *
 * Full PIN auth module with rotary entry, lockout, and dual identity modes.
 *       Dual wired/wired identity modes backed by non-volatile failure
 *       limits, roll-count PIN security, and automatic device state
 *       enforcement.  Wired mode enables USB composite gadget via
 *       enable_vault_usb.sh; wireless mode tears down USB and brings
 *       up BLE via enable_vault_ble.sh.
 *
 *       Manifest ref: Asset IV -- Crypto Security Token Vault Gateway
 */
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <stdint.h>
#include <string.h>
#include <errno.h>
#include <linux/input.h>
#include "h2_ui.h"
#include "log_manager.h"

#define STATE_FILE  "/data/vault/failures.dat"
#define MAX_FAILURES 3
#define LOCKOUT_SEC  300

/* PIN: 4 digits, default 4-2-9-1 as per manifest */
static int master_pin[4] = {4, 2, 9, 1};
static int entered_pin[4] = {0, 0, 0, 0};
static int current_digit_idx = 0;

typedef enum { MODE_NONE, MODE_WIRED, MODE_WIRELESS } VaultMode;
static VaultMode current_mode = MODE_NONE;

/* UI objects */
static lv_obj_t *pin_label = NULL;
static lv_obj_t *warning_label = NULL;
static lv_obj_t *digit_labels[4] = {NULL};
static lv_obj_t *cursor_label = NULL;
static lv_obj_t *status_label = NULL;
static lv_obj_t *mode_label = NULL;

static int get_failure_count(void) {
    FILE *f = fopen(STATE_FILE, "r");
    if (!f) return 0;
    int count = 0;
    if (fscanf(f, "%d", &count) <= 0) count = 0;
    fclose(f);
    return count;
}

static void set_failure_count(int count) {
    FILE *f = fopen(STATE_FILE, "w");
    if (f) { fprintf(f, "%d", count); fclose(f); }
}

static int is_usb_plugged_in(void) {
    int fd = open("/sys/class/power_supply/usb/online", O_RDONLY);
    if (fd == -1) return 0;
    char status = '0';
    if (read(fd, &status, 1) <= 0) status = '0';
    close(fd);
    return (status == '1');
}

static void render_pin_screen(int attempts_left) {
    char buf[64];

    /* Warning */
    snprintf(buf, sizeof(buf), "Attempts before lockdown: %d", attempts_left);
    lv_label_set_text(warning_label, buf);
    lv_obj_set_style_text_color(warning_label, lv_color_make(253, 32, 32), LV_PART_MAIN);

    /* Digit boxes */
    for (int i = 0; i < 4; i++) {
        char digit_str[16];
        snprintf(digit_str, sizeof(digit_str), "%d", entered_pin[i]);
        lv_label_set_text(digit_labels[i], digit_str);

        if (i == current_digit_idx) {
            lv_obj_set_style_text_color(digit_labels[i], COLOR_ACCENT, LV_PART_MAIN);
            lv_obj_set_style_border_color(digit_labels[i],
                lv_color_make(0, 224, 0), LV_PART_MAIN);
            lv_obj_set_style_border_width(digit_labels[i], 2, LV_PART_MAIN);
        } else {
            lv_obj_set_style_text_color(digit_labels[i], COLOR_PRIMARY, LV_PART_MAIN);
            lv_obj_set_style_border_width(digit_labels[i], 1, LV_PART_MAIN);
        }
    }
}

static int enforce_pin_authorization(int input_fd) {
    int failures = get_failure_count();

    /* Lockout check */
    if (failures >= MAX_FAILURES) {
        char log_msg[128];
        snprintf(log_msg, sizeof(log_msg),
            "PIN lockout triggered. %d failures recorded. Enforcing %ds cooldown.",
            failures, LOCKOUT_SEC);
        LOG_ERR("%s", log_msg);

        for (int penalty_sec = LOCKOUT_SEC; penalty_sec > 0; penalty_sec--) {
            char countdown[64];
            snprintf(countdown, sizeof(countdown),
                "DEVICE LOCKED DOWN\nRetry in: %d s", penalty_sec);
            lv_label_set_text(status_label, countdown);
            lv_obj_set_style_text_color(status_label,
                lv_color_make(255, 0, 0), LV_PART_MAIN);
            lv_timer_handler();

            /* Check for back key during lockout to allow early exit */
            struct input_event ev;
            if (input_fd >= 0) {
                int flags = fcntl(input_fd, F_GETFL, 0);
                fcntl(input_fd, F_SETFL, flags | O_NONBLOCK);
                if (read(input_fd, &ev, sizeof(struct input_event)) > 0) {
                    if (ev.type == EV_KEY && ev.code == H2_KEY_BACK && ev.value == 1) {
                        fcntl(input_fd, F_SETFL, flags);
                        return -1;
                    }
                }
                fcntl(input_fd, F_SETFL, flags);
            }
            sleep(1);
        }
        set_failure_count(0);
        failures = 0;
    }

    render_pin_screen(MAX_FAILURES - failures);

    /* PIN entry loop */
    struct input_event ev;
    while (current_digit_idx < 4) {
        lv_timer_handler();
        if (input_fd >= 0 && read(input_fd, &ev, sizeof(struct input_event)) > 0) {
            if (ev.type == EV_REL && ev.code == H2_ROTARY_REL) {
                if (ev.value > 0)
                    entered_pin[current_digit_idx] = (entered_pin[current_digit_idx] + 1) % 10;
                else
                    entered_pin[current_digit_idx] = (entered_pin[current_digit_idx] - 1 + 10) % 10;
                render_pin_screen(MAX_FAILURES - failures);
            } else if (ev.type == EV_KEY && ev.code == H2_KEY_PLAY && ev.value == 1) {
                current_digit_idx++;
                if (current_digit_idx < 4) render_pin_screen(MAX_FAILURES - failures);
            } else if (ev.type == EV_KEY && ev.code == H2_KEY_BACK && ev.value == 1) {
                if (current_digit_idx > 0) {
                    current_digit_idx--;
                    render_pin_screen(MAX_FAILURES - failures);
                }
            }
        }
        usleep(5000);
    }

    /* Verify PIN */
    if (memcmp(master_pin, entered_pin, sizeof(master_pin)) == 0) {
        set_failure_count(0);
        LOG_INF("PIN verified. Authorization granted.");
        lv_label_set_text(status_label, "ACCESS GRANTED\nKEY INJECTED");
        lv_obj_set_style_text_color(status_label, COLOR_ACCENT, LV_PART_MAIN);
        lv_timer_handler();
        sleep(2);
        return 0;
    } else {
        failures++;
        set_failure_count(failures);
        char log_msg[128];
        snprintf(log_msg, sizeof(log_msg),
            "Invalid PIN attempt logged. Failure level: %d/%d",
            failures, MAX_FAILURES);
        LOG_ERR("%s", log_msg);

        lv_label_set_text(status_label, "INVALID PIN\nATTEMPT LOGGED");
        lv_obj_set_style_text_color(status_label,
            lv_color_make(255, 0, 0), LV_PART_MAIN);
        lv_timer_handler();
        sleep(2);
        return -1;
    }
}

int main(void) {
    init_h2_graphics_runtime("CRYPTO VAULT GUARD");
    LOG_INF("vault module started (full PIN auth mode)");

    lv_obj_t *scr = lv_scr_act();

    /* Title */
    lv_obj_t *title = lv_label_create(scr);
    lv_label_set_text(title, "SECURITY TOKEN VAULT");
    lv_obj_align(title, LV_ALIGN_TOP_MID, 0, 10);
    lv_obj_set_style_text_color(title, COLOR_PRIMARY, LV_PART_MAIN);
    lv_obj_set_style_text_font(title, &lv_font_montserrat_14, LV_PART_MAIN);

    /* PIN entry card */
    lv_obj_t *pin_card = lv_obj_create(scr);
    lv_obj_set_size(pin_card, 280, 140);
    lv_obj_align(pin_card, LV_ALIGN_CENTER, 0, -5);

    lv_obj_t *pin_title = lv_label_create(pin_card);
    lv_label_set_text(pin_title, "ENTER PIN");
    lv_obj_align(pin_title, LV_ALIGN_TOP_MID, 0, 5);
    lv_obj_set_style_text_color(pin_title, COLOR_MUTED, LV_PART_MAIN);

    warning_label = lv_label_create(pin_card);
    lv_obj_align(warning_label, LV_ALIGN_TOP_MID, 0, 22);
    lv_obj_set_style_text_font(warning_label, &lv_font_montserrat_10, LV_PART_MAIN);

    /* Four digit display boxes */
    for (int i = 0; i < 4; i++) {
        digit_labels[i] = lv_label_create(pin_card);
        lv_label_set_text(digit_labels[i], "0");
        lv_obj_set_style_text_font(digit_labels[i], &lv_font_montserrat_22, LV_PART_MAIN);
        lv_obj_set_style_border_side(digit_labels[i],
            LV_BORDER_SIDE_FULL, LV_PART_MAIN);
        lv_obj_set_style_border_color(digit_labels[i], COLOR_MUTED, LV_PART_MAIN);
        lv_obj_set_style_border_width(digit_labels[i], 1, LV_PART_MAIN);
        lv_obj_set_style_pad_all(digit_labels[i], 8, LV_PART_MAIN);
        lv_obj_align(digit_labels[i], LV_ALIGN_CENTER, -60 + (i * 40), -5);
    }

    /* Status area (used for lockout / result) */
    status_label = lv_label_create(scr);
    lv_label_set_text(status_label, "");
    lv_obj_set_style_text_font(status_label, &lv_font_montserrat_12, LV_PART_MAIN);
    lv_obj_align(status_label, LV_ALIGN_BOTTOM_MID, 0, -30);

    /* Mode indicator */
    mode_label = lv_label_create(scr);
    lv_obj_align(mode_label, LV_ALIGN_BOTTOM_LEFT, 8, -8);
    lv_obj_set_style_text_color(mode_label, COLOR_MUTED, LV_PART_MAIN);
    lv_obj_set_style_text_font(mode_label, &lv_font_montserrat_10, LV_PART_MAIN);
    lv_label_set_text(mode_label, "[WHEEL] Digit  [PLAY] Next  [BACK] Exit");

    /* Open input */
    int input_fd = open("/dev/input/event0", O_RDONLY | O_NONBLOCK);
    if (input_fd < 0) {
        LOG_ERR("cannot open input device");
        return 1;
    }

    /* Run PIN authorization */
    int auth_result = enforce_pin_authorization(input_fd);
    if (auth_result != 0) {
        close(input_fd);
        LOG_INF("vault module exited (auth failed)");
        return 0;
    }

    /* Authorized: enter runtime mode loop */
    while (1) {
        lv_timer_handler();

        int usb_active = is_usb_plugged_in();

        if (usb_active) {
            if (current_mode != MODE_WIRED) {
                LOG_INF("switching to WIRED mode");
                system("hciconfig hci0 down 2>/dev/null");
                system("/usr/bin/enable_vault_usb.sh 2>/dev/null");
                current_mode = MODE_WIRED;
            }
            lv_label_set_text(status_label, "MODE: SECURE WIRED SMARTCARD\nUSB Token: Operational (CCID)");
            lv_obj_set_style_text_color(status_label,
                lv_color_make(0, 68, 136), LV_PART_MAIN);
        } else {
            if (current_mode != MODE_WIRELESS) {
                LOG_INF("switching to WIRELESS BLE mode");
                system("echo \"\" > /sys/kernel/config/usb_gadget/wh_tool/UDC 2>/dev/null");
                system("/usr/bin/enable_vault_ble.sh 2>/dev/null");
                current_mode = MODE_WIRELESS;
            }
            lv_label_set_text(status_label, "MODE: WIRELESS BLE SMARTCARD\nBLE Identity: Active");
            lv_obj_set_style_text_color(status_label,
                lv_color_make(0, 204, 68), LV_PART_MAIN);
        }

        /* Check for BACK to exit and purge */
        struct input_event ev;
        if (input_fd >= 0 && read(input_fd, &ev, sizeof(struct input_event)) > 0) {
            if (ev.type == EV_KEY && ev.code == H2_KEY_BACK && ev.value == 1) {
                lv_label_set_text(status_label, "Purging keys from RAM... Locking.");
                lv_obj_set_style_text_color(status_label,
                    lv_color_make(255, 0, 0), LV_PART_MAIN);
                lv_timer_handler();
                LOG_INF("vault key purge triggered by user");
                sleep(1);
                break;
            }
        }
        usleep(200000);
    }

    close(input_fd);
    LOG_INF("vault module exited");
    return 0;
}