/* Motenpoche 
 *
 * (c) 2023 Daniele Lacamera <root@danielinux.net>
 *
 *
 * Motenpoche is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * Motenpoche is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1335, USA
 *
 */
#include <stdint.h>
#include "ui.h"
#include "display.h"
#include "bsp/board.h"
#include "hardware/gpio.h"
#include "cryptoengine.h"
#include "wolfssl/wolfcrypt/settings.h"
#include "wolfssl/wolfcrypt/misc.h"
#include "fsm.h"
#include "display.h"
#include "flash.h"
#include "hid.h"


#define PASSPHRASE_MAX 128

#ifndef GREEN_LED
#define GREEN_LED 16
#endif


static const char welcome_message0[]= "  There is no  ";
static const char welcome_message1[]= "knowledge that ";
static const char welcome_message2[]= " is not power. ";

extern void system_reboot(void);
static uint16_t service_list_selected = 0;
static uint16_t service_list_count = 0;
static void ui_service_list(uint32_t count);



/* UI current state globals */
static int current_row = 1;
static int current_char_pos = 0;
static char current_char = ' ';
static char passphrase[PASSPHRASE_MAX];
static int blinking = 0;

void ui_event_cb(enum vault_state old_st, enum vault_state st)
{
    int result = 0;
    int ret = -1;
    current_char = ' ';
    (void)old_st;
    switch (st) {
        case VAULT_BOOTUP:
            display_clear();
            display_text(0, "Type Passphrase");
            ForceZero(passphrase, PASSPHRASE_MAX);
            current_char_pos = 0;
            break;
        case VAULT_VERIFY_PASSPHRASE:
            char msg[17];
            display_text(0, "Authenticating..");
            display_text(1, "                ");
            display_text(2, "                ");
            display_text(3,"                 ");
            ret = cryptoengine_verify_passphrase(passphrase, &result);
            if ((ret != 0) || (result != 1)) {
                fsm_set(VAULT_VERIFY_FAILED);
                return;
            }
            gpio_put(GREEN_LED, 0);
            fsm_set(VAULT_MAIN_MENU);
            return;
            break;
        case VAULT_MAIN_MENU:
            display_text(0,"                 ");
            display_text(1,"                 ");
            display_text(2,"                 ");
            display_text(3,"                 ");
            display_text_inverse(0, 2, "[Main Menu]");
            break;
        case VAULT_SETTINGS_MENU:
            display_text(0,"                 ");
            display_text(1,"                 ");
            display_text(2,"                 ");
            display_text(3,"                 ");
            display_text_inverse(0, 3, "[Settings]");
            break;
        case VAULT_SERVICE_LIST:
            uint16_t count, spot;
            cryptoengine_service_count(&count, &spot);
            service_list_count = count;
            display_text(0,"                 ");
            display_text(1,"                 ");
            display_text(2,"                 ");
            display_text(3,"                 ");
            display_text_inverse(0, 3, "[Services]");
            break;
    }
    ui_task();
}


static const char main_menu_entry[3][17] = {
    "[~  Services  ~]",
    "[~  Settings  ~]",
    "[~  Reboot    ~]"
};


#define MAINMENU_ENTRIES 3

static const char settings_menu_entry[3][17] = {
    "[~ Keyb map   ~]",
    "[~ Paste mode ~]",
    "[~ Forget All ~]"
};
#define SETTINGSMENU_ENTRIES 3

static int menu_sel = 0;
static int old_menu_sel = 0xFF;

static const char Keymaps[][17] = {
    "      [US]",
    "      [FR]",
    "      [DE]",
    "      [UK]",
    "      [IT]"
};

static int cur_keymap = 0;

#define KEYMAPS_ENTRIES 5


#define PASTEMODE_PASS_ONLY 0
#define PASTEMODE_PASS_ENTER 1
#define PASTEMODE_USER_TAB_PASS_ENTER 2
#define PASTEMODE_USER_ENTER_PASS_ENTER 3

static const char PasteModes[][17] = {
    "[Password only ]",
    "[Password+Enter]",
    "[UserTabPswdEnt]",
    "[UserEntPswdEnt]"
};

static int cur_pastemode = 0;

#define PASTEMODE_ENTRIES 4

static const char Are_you_sure[][17] = {
    "Sure? NO    ",
    "Sure? Cancel",
    "Sure? My mistake",
    "Sure? Whut?",
    "Sure? YES DO IT!"
};

static int cur_wipe_choice = 0;
#define AREYOUSURE_ENTRIES 5



static void mainmenu_action(void)
{
    switch (menu_sel) {
        case 0: /* Services */
            fsm_set(VAULT_SERVICE_LIST);
            break;
        case 1: /* Settings */
            fsm_set(VAULT_SETTINGS_MENU);
            break;
        case 2: /* Reboot */
            system_reboot();
            break;
    }
}

static void settingsmenu_action(void)
{
    switch (menu_sel) {
        case 0: /* Keymap */
            cur_keymap++;
            if (cur_keymap == KEYMAPS_ENTRIES)
                cur_keymap = 0;
            break;
        case 1: /* Paste mode */
            cur_pastemode++;
            if (cur_pastemode == PASTEMODE_ENTRIES)
                cur_pastemode = 0;
            break;
        case 2: /* Forget */
            cur_wipe_choice++;
            if (cur_wipe_choice == AREYOUSURE_ENTRIES)
                cur_wipe_choice = 0;
            break;
    }
    /* Force redraw */
    old_menu_sel = 0xFF;
}

static void settingsmenu_confirm(void)
{
    int i;
    if ((menu_sel == 2) && (cur_wipe_choice == 4)) {
        for (i = 0; i < VAULT_FLASH_SIZE; i+= SPI_FLASH_SECTOR_SIZE)
            flash_sector_erase(i);
        system_reboot();
    } else
        settingsmenu_action();
}


static void ui_service_list(uint32_t count)
{
    static int old_selected = -1;
    int i = 0;
    uint32_t address = SVC_FLASH_OFFSET;
    uint32_t a_sel, a_prev = 0, a_next = 0;
    struct vault_service svc;

    if (old_selected == (int)service_list_selected)
        return;
    old_selected = service_list_selected;

    if (count == 0)
        return;


    flash_read(address + service_list_selected * SVC_SIZE,
            (void *)&svc, SVC_SIZE);
    if (svc.flags != SVC_FLAG_ACTIVE) {
        i =  0;
        while (address < VAULT_FLASH_SIZE) {
            flash_read(address + i * SVC_SIZE, &svc, SVC_SIZE);
            if (svc.flags == SVC_FLAG_ACTIVE) {
                service_list_selected = i;
                break;
            }
            i += SVC_SIZE;
        }
    }
    a_sel = SVC_FLASH_OFFSET + service_list_selected * SVC_SIZE;
    a_next = 0;
    a_prev = 0;

    /* Find previous entry */
    if (service_list_selected != 0) {
        i = 1;
        address = SVC_FLASH_OFFSET + SVC_SIZE * (service_list_selected - i);
        while (address >= SVC_FLASH_OFFSET) {
            flash_read(address, (void *)&svc, SVC_SIZE);
            if (svc.flags == SVC_FLAG_ACTIVE) {
                a_prev = address;
                break;
            }
            i++;
            address = SVC_FLASH_OFFSET + SVC_SIZE * (service_list_selected - i);
        }
    }
    /* Find next entry */
    i = 1;
    address = SVC_FLASH_OFFSET + SVC_SIZE * (service_list_selected + i);
    while (address < VAULT_FLASH_SIZE) {
        flash_read(address, (void *)&svc, SVC_SIZE);
        if (svc.flags == SVC_FLAG_ACTIVE) {
            a_next = address;
            break;
        } else if (svc.flags == SVC_FLAG_UNUSED) {
            break;
        }
        i++;
        address = SVC_FLASH_OFFSET + SVC_SIZE * (service_list_selected + i);
    }

    display_text(1,"                 ");
    display_text(2,"                 ");
    display_text(3,"                 ");

    if (a_prev != 0) {
        flash_decrypt_read_svc(&svc, a_prev);
        display_text(1, svc.name);
    }
    if (a_next != 0) {
        flash_decrypt_read_svc(&svc, a_next);
        display_text(3, svc.name);
    }
    flash_decrypt_read_svc(&svc, a_sel);
    display_text_inverse(2, 0, svc.name);
}


static void paste_password(void)
{
    struct vault_service svc;
    uint32_t address;
    if (fsm_get() != VAULT_SERVICE_LIST)
        return;
    if (flash_decrypt_read_svc(&svc, SVC_FLASH_OFFSET +
                service_list_selected * SVC_SIZE) < 0)
        return;

    switch(cur_pastemode) {
        case PASTEMODE_PASS_ONLY:
            hid_keys_string_send(svc.pass);
            break;
        case PASTEMODE_PASS_ENTER:
            hid_keys_string_send(svc.pass);
            hid_keys_string_send("\r");
            break;
        case PASTEMODE_USER_TAB_PASS_ENTER:
            hid_keys_string_send(svc.user);
            hid_keys_string_send("\t");
            hid_keys_string_send(svc.pass);
            hid_keys_string_send("\r");
            break;
        case PASTEMODE_USER_ENTER_PASS_ENTER:
            hid_keys_string_send(svc.user);
            hid_keys_string_send("\t");
            hid_keys_string_send(svc.pass);
            hid_keys_string_send("\r");
            break;
    }
}

void ui_task(void)
{
    const uint32_t interval_ms = 500;
    static uint32_t start_ms = 0;

    if ( board_millis() - start_ms >= interval_ms) {
        blinking ^= 1;
        start_ms = board_millis();
    }

    switch(fsm_get()) {
        case VAULT_TOFU:
            gpio_put(17, blinking ^ 1);
            display_text(0, "    Welcome!");
            display_text(1, "Run tools on PC ");
            display_text(2, " to initialize  ");
            break;
        case VAULT_BOOTUP:
            char pass_prompt[16];
            int i;
            int sz;
            gpio_put(17, blinking ^ 1);
            sz = current_char_pos;
            if (sz > 15)
                sz = 15;
            memset(pass_prompt, 0, 16);
            for (i = 0; i < sz; i++) {
               pass_prompt[i] = '*';
            }
            if (!blinking)
                pass_prompt[sz] = current_char; 
            display_text(1, pass_prompt);
            if (blinking) {
                char cursor[2] = {current_char, 0};
                display_text_inverse(1, sz, cursor);
            }
            break;
        case VAULT_VERIFY_PASSPHRASE:
            display_text(2, "Verifying...");
            gpio_put(17, blinking ^ 1);
            break;
        case VAULT_VERIFY_FAILED:
            display_text(0, " Bad Passphrase ");
            gpio_put(17, 0);
            if (blinking) {
                display_text(1,"    ");
                display_text(2,"   ");
                display_text_inverse(1, 4, "(x_x)");
                display_text_inverse(2, 3, "Try again");
            } else {
                display_text(1, "    (x_x)   ");
                display_text(2, "            ");
            }
            break;
        case VAULT_MAIN_MENU:
            display_text(2,main_menu_entry[menu_sel]);
            break;
        case VAULT_SETTINGS_MENU:

            if (old_menu_sel != menu_sel) {
                old_menu_sel = menu_sel;

                display_text(1, "                ");
                display_text(2, "                ");
                display_text(3, "                ");
                display_text(2,settings_menu_entry[menu_sel]);
                switch(menu_sel) {
                    case 0: /* Keymap */
                        display_text(3, Keymaps[cur_keymap]);
                        break;
                    case 1:
                        display_text(3, PasteModes[cur_pastemode]);
                        break;
                    case 2:
                        display_text(3, Are_you_sure[cur_wipe_choice]);
                        break;
                }
            }
            break;
        case VAULT_SERVICE_LIST:
            if (service_list_count == 0) {
                display_text(2, "  No Services  "); 
                break;
            } else {
                ui_service_list(service_list_count);
            }
            break;
    }
}

void ui_key_button(void)
{
    static int alter = 0;
    enum vault_state vs = fsm_get();
    if (vs == VAULT_BOOTUP) {
        passphrase[current_char_pos] = current_char;
        current_char_pos++;
        if(current_char_pos >= PASSPHRASE_MAX) {
            fsm_set(VAULT_VERIFY_FAILED);
            return;
        }
        passphrase[current_char_pos] = 0;
        gpio_put(17, alter);
        alter ^=1;
    } else if (vs == VAULT_MAIN_MENU)
        mainmenu_action();
    else if (vs == VAULT_SETTINGS_MENU)
        settingsmenu_action();
    else if (vs == VAULT_SERVICE_LIST)
        paste_password();
}

void ui_back_button(void)
{
    enum vault_state vault_state = fsm_get();
    if (vault_state == VAULT_BOOTUP) {
        if (current_char_pos > 0) {
            passphrase[current_char_pos] = '\0';
            current_char_pos--;
        }
        display_text(1, "               ");
    }
    else if (vault_state == VAULT_VERIFY_FAILED) {
        current_char_pos = 0;
        fsm_set(VAULT_BOOTUP);
    }
    else if (vault_state == VAULT_SETTINGS_MENU)
        fsm_set(VAULT_MAIN_MENU);
    else if (vault_state == VAULT_SERVICE_LIST)
        fsm_set(VAULT_MAIN_MENU);
}

void ui_confirm_button(void)
{
    enum vault_state vault_state = fsm_get();
    if (vault_state == VAULT_BOOTUP) {
        current_char_pos = 0;
        fsm_set(VAULT_VERIFY_PASSPHRASE);
    }
    else if (vault_state == VAULT_MAIN_MENU)
        mainmenu_action();
    else if (vault_state == VAULT_SETTINGS_MENU)
        settingsmenu_confirm();
    else if (vault_state == VAULT_SERVICE_LIST)
        paste_password();
}

void ui_rot_up(void)
{
    enum vault_state vault_state = fsm_get();
    if (vault_state == VAULT_BOOTUP) {
        current_char++;
        if (current_char > 126)
            current_char = ' ';
    }
    else if (vault_state == VAULT_MAIN_MENU) {
        menu_sel++;
        if (menu_sel == MAINMENU_ENTRIES)
            menu_sel = 0;
    }
    else if (vault_state == VAULT_SETTINGS_MENU) {
        menu_sel++;
        cur_wipe_choice = 0;
        if (menu_sel == MAINMENU_ENTRIES)
            menu_sel = 0;
    }
    else if (vault_state == VAULT_SERVICE_LIST) {
        if (service_list_selected > 0)
            service_list_selected --;
    }
}

void ui_rot_down(void)
{
    enum vault_state vault_state = fsm_get();
    if (vault_state == VAULT_BOOTUP) {
        current_char--;
        if (current_char < ' ')
            current_char = 126;
    }
    else if (vault_state == VAULT_MAIN_MENU) {
        menu_sel--;
        if (menu_sel < 0)
            menu_sel = MAINMENU_ENTRIES - 1;
    }
    else if (vault_state == VAULT_SETTINGS_MENU) {
        menu_sel--;
        cur_wipe_choice = 0;
        if (menu_sel < 0)
            menu_sel = SETTINGSMENU_ENTRIES - 1;
    }
    else if (vault_state == VAULT_SERVICE_LIST) {
        struct vault_service svc;
        int next = service_list_selected + 1;
        while (1) { 
            flash_read(SVC_FLASH_OFFSET + next * SVC_SIZE, &svc, SVC_SIZE);
            if (svc.flags == SVC_FLAG_UNUSED)
                return; /* Keep old value */
            else if (svc.flags == SVC_FLAG_ERASED) {
                next++;
            } else {
                service_list_selected = next;
                break;
            }
        }
    }
}

void ui_init(i2c_inst_t *i2c)
{
    int i;
    display_init(i2c);
    display_clear();
#ifdef SHOW_MOTD
    display_scroll(0x3f);
    display_text(0, welcome_message0);
    display_text(1, welcome_message1);
    display_text(2, welcome_message2);
    for (i = 0x3f; i >= 0x20; i--) {
        display_scroll(i);
        sleep_us(50000);
    }
    for (i = display_getcontrast(); i >= 0; i--) {
        display_setcontrast(i);
        sleep_us(10000);
    }
    display_scroll(0);
    display_clear();
#endif
    display_setcontrast(0xcf);
    current_char_pos = 0; 
    sleep_us(50000);
    display_scroll(0x20);


    if (cryptoengine_check_vault() < 0)
        fsm_set(VAULT_TOFU);
    else
        fsm_set(VAULT_BOOTUP);
}

void ui_event(uint32_t ev)
{
}
