#include "ui.h"

#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>

#include "buttons.h"
#include "display.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "mnemonic.h"
#include "secure_mem.h"

typedef enum {
    UI_MAIN_GENERATE = 0,
    UI_MAIN_ENTER,
    UI_MAIN_VIEW,
} ui_main_menu_item_t;

#define UI_CHAR_W_SMALL 6
#define UI_CHAR_W_LARGE 12

typedef struct {
    size_t count;
    size_t first_idx[3];
    int exact_idx;
} prefix_matches_t;

typedef struct {
    char phrase[256];
    bool has_phrase;

    ui_main_menu_item_t main_sel;

    int letter_idx; /* 0..25 */
    char prefix[12];
    size_t prefix_len;

    char entered_words[12][12];
    size_t entered_count;
} ui_state_t;

static btn_event_t wait_any_button(void)
{
    btn_event_t evt;
    do {
        evt = buttons_wait(0xFFFFFFFFu);
    } while (evt.kind == BTN_EVT_NONE);
    return evt;
}

static void draw_splash(void)
{
    ESP_ERROR_CHECK(display_clear(RGB565_BLACK));
    ESP_ERROR_CHECK(display_draw_text_centered(60, "N_SIGNER", RGB565_RED, RGB565_BLACK));
}

static void draw_main_menu_row(bool has_phrase, int row_idx, bool selected)
{
    static const char *items[3] = {"Generate", "Enter", "View"};
    static const int ys[3] = {28, 52, 76};
    bool disabled = (row_idx == UI_MAIN_VIEW && !has_phrase);
    char line[24];

    snprintf(line, sizeof(line), "%s%s", selected ? "> " : "  ", items[row_idx]);
    ESP_ERROR_CHECK(display_fill_rect(24, ys[row_idx], 200, 16, RGB565_BLACK));
    ESP_ERROR_CHECK(display_draw_text(
        24,
        ys[row_idx],
        line,
        disabled ? RGB565_DIM_GRAY : (selected ? RGB565_RED : RGB565_WHITE),
        RGB565_BLACK));
}

static void draw_main_menu_full(const ui_state_t *st)
{
    ESP_ERROR_CHECK(display_clear(RGB565_BLACK));
    ESP_ERROR_CHECK(display_draw_text_centered(4, "MAIN MENU", RGB565_WHITE, RGB565_BLACK));

    for (int i = 0; i < 3; ++i) {
        draw_main_menu_row(st->has_phrase, i, ((int)st->main_sel == i));
    }

    ESP_ERROR_CHECK(display_draw_text_small(12, 120, "D2 up   D1 dn   D0 ok", RGB565_DIM_GRAY, RGB565_BLACK));
}

static void draw_main_menu_selection_change(bool has_phrase, int prev_sel, int new_sel)
{
    draw_main_menu_row(has_phrase, prev_sel, false);
    draw_main_menu_row(has_phrase, new_sel, true);
}

static void split_phrase_words(const char *phrase, char words[12][12])
{
    char buf[256];
    char *tok;
    size_t i = 0;

    memset(words, 0, 12 * 12);

    if (phrase == NULL) {
        return;
    }

    strncpy(buf, phrase, sizeof(buf) - 1);
    buf[sizeof(buf) - 1] = '\0';

    tok = strtok(buf, " \t\r\n");
    while (tok != NULL && i < 12) {
        strncpy(words[i], tok, 11);
        words[i][11] = '\0';
        i++;
        tok = strtok(NULL, " \t\r\n");
    }
}

static void draw_word_grid(const char words[12][12], int highlight_idx)
{
    const int col_w = 80;
    const int row_h = 24;

    for (int i = 0; i < 12; ++i) {
        int row = i / 3;
        int col = i % 3;
        int x = col * col_w + 2;
        int y = 14 + row * row_h;
        char line[20];

        snprintf(line, sizeof(line), "%2d. %.11s", i + 1, words[i]);
        ESP_ERROR_CHECK(display_draw_text_small(
            x,
            y,
            line,
            (i == highlight_idx) ? RGB565_RED : RGB565_WHITE,
            RGB565_BLACK));
    }
}

static void draw_view_page(const ui_state_t *st)
{
    char words[12][12];

    split_phrase_words(st->phrase, words);

    ESP_ERROR_CHECK(display_clear(RGB565_BLACK));
    ESP_ERROR_CHECK(display_draw_text_small_centered(2, "YOUR MNEMONIC", RGB565_WHITE, RGB565_BLACK));
    draw_word_grid(words, -1);
    ESP_ERROR_CHECK(display_draw_text_small_centered(120, "any key = back", RGB565_DIM_GRAY, RGB565_BLACK));
}

static void draw_generate_page(void)
{
    ESP_ERROR_CHECK(display_clear(RGB565_BLACK));
    ESP_ERROR_CHECK(display_draw_text_centered(56, "GENERATING...", RGB565_WHITE, RGB565_BLACK));
}

static void clear_enter_words(ui_state_t *st)
{
    secure_memzero(st->entered_words, sizeof(st->entered_words));
    st->entered_count = 0;
    st->prefix_len = 0;
    st->prefix[0] = '\0';
}

static void letter_cell_pos(int idx, int *x, int *y)
{
    int row = (idx < 13) ? 0 : 1;
    int col = (idx < 13) ? idx : (idx - 13);
    *x = 10 + col * 16;
    *y = row ? 68 : 44;
}

static void compute_allowed_next_letters(const ui_state_t *st, bool allowed[26])
{
    for (int i = 0; i < 26; ++i) {
        allowed[i] = false;
    }

    if (st->prefix_len == 0) {
        for (int i = 0; i < 26; ++i) {
            allowed[i] = true;
        }
        return;
    }

    for (size_t i = 0; i < 2048; ++i) {
        const char *w = mnemonic_word_at(i);
        size_t wlen;
        char c;
        int idx;

        if (w == NULL) {
            continue;
        }
        if (strncmp(w, st->prefix, st->prefix_len) != 0) {
            continue;
        }

        wlen = strlen(w);
        if (wlen <= st->prefix_len) {
            continue;
        }

        c = w[st->prefix_len];
        if (c < 'a' || c > 'z') {
            continue;
        }

        idx = (int)(c - 'a');
        allowed[idx] = true;
    }
}

static int find_next_allowed_letter(int current, int delta, const bool allowed[26])
{
    int idx = current;

    for (int i = 0; i < 26; ++i) {
        idx = (idx + delta + 26) % 26;
        if (allowed[idx]) {
            return idx;
        }
    }

    return current;
}

static int first_allowed_letter(const bool allowed[26])
{
    for (int i = 0; i < 26; ++i) {
        if (allowed[i]) {
            return i;
        }
    }
    return 0;
}

static void draw_enter_letter_cell(int idx, bool selected, const bool allowed[26])
{
    int x;
    int y;
    char c[2] = {(char)('A' + idx), '\0'};
    uint16_t color = RGB565_WHITE;

    letter_cell_pos(idx, &x, &y);
    ESP_ERROR_CHECK(display_fill_rect(x, y, 12, 16, RGB565_BLACK));

    if (!allowed[idx]) {
        color = RGB565_DIM_GRAY;
    }
    if (selected && allowed[idx]) {
        color = RGB565_RED;
    }

    ESP_ERROR_CHECK(display_draw_text(x, y, c, color, RGB565_BLACK));
}

static void draw_enter_letter_cursor_change(int prev_idx, int new_idx, const bool allowed[26])
{
    draw_enter_letter_cell(prev_idx, false, allowed);
    draw_enter_letter_cell(new_idx, true, allowed);
}

static void compute_prefix_matches(const ui_state_t *st, prefix_matches_t *m)
{
    m->count = 0;
    m->first_idx[0] = 0;
    m->first_idx[1] = 0;
    m->first_idx[2] = 0;
    m->exact_idx = -1;

    if (st->prefix_len == 0) {
        return;
    }

    for (size_t i = 0; i < 2048; ++i) {
        const char *w = mnemonic_word_at(i);
        if (w == NULL) {
            continue;
        }

        if (strncmp(w, st->prefix, st->prefix_len) == 0) {
            if (m->count < 3) {
                m->first_idx[m->count] = i;
            }
            m->count++;
            if (strcmp(w, st->prefix) == 0 && m->exact_idx < 0) {
                m->exact_idx = (int)i;
            }
        }
    }
}

static void draw_enter_compact_static(const ui_state_t *st, const bool allowed[26])
{
    char hdr[24];
    unsigned current_word = (unsigned)(st->entered_count + 1);

    ESP_ERROR_CHECK(display_clear(RGB565_BLACK));

    if (current_word > 12) {
        current_word = 12;
    }
    snprintf(hdr, sizeof(hdr), "WORD %u/12", current_word);
    ESP_ERROR_CHECK(display_draw_text(10, 8, hdr, RGB565_WHITE, RGB565_BLACK));

    for (int i = 0; i < 26; ++i) {
        draw_enter_letter_cell(i, i == st->letter_idx, allowed);
    }
}

static int draw_word_with_prefix_large(int x, int y, const char *word, size_t prefix_len)
{
    char pref[12];
    char rest[12];
    size_t wlen;
    size_t plen;

    if (word == NULL) {
        return x;
    }

    wlen = strlen(word);
    plen = (prefix_len > wlen) ? wlen : prefix_len;

    memset(pref, 0, sizeof(pref));
    memset(rest, 0, sizeof(rest));

    if (plen > 0) {
        memcpy(pref, word, plen);
    }
    if (wlen > plen) {
        memcpy(rest, word + plen, wlen - plen);
    }

    ESP_ERROR_CHECK(display_draw_text(x, y, pref, RGB565_RED, RGB565_BLACK));
    ESP_ERROR_CHECK(display_draw_text(x + ((int)plen * UI_CHAR_W_LARGE), y, rest, RGB565_WHITE, RGB565_BLACK));

    return x + ((int)wlen * UI_CHAR_W_LARGE);
}

static void draw_match_row_large(const ui_state_t *st)
{
    int x = 4;

    for (size_t i = 0; i < 2048; ++i) {
        const char *w = mnemonic_word_at(i);
        size_t wlen;
        int needed;

        if (w == NULL) {
            continue;
        }
        if (strncmp(w, st->prefix, st->prefix_len) != 0) {
            continue;
        }

        wlen = strlen(w);
        needed = (int)(wlen * UI_CHAR_W_LARGE) + UI_CHAR_W_LARGE;
        if (x + needed > DISPLAY_WIDTH) {
            break;
        }

        x = draw_word_with_prefix_large(x, 88, w, st->prefix_len);
        ESP_ERROR_CHECK(display_draw_text(x, 88, " ", RGB565_WHITE, RGB565_BLACK));
        x += UI_CHAR_W_LARGE;
    }
}

static void draw_enter_compact_dynamic(const ui_state_t *st, const prefix_matches_t *m)
{
    ESP_ERROR_CHECK(display_fill_rect(0, 86, DISPLAY_WIDTH, 49, RGB565_BLACK));

    if (st->prefix_len == 0) {
        ESP_ERROR_CHECK(display_draw_text_small(4, 118, "D0=letter  D1h=back  D2h=del", RGB565_DIM_GRAY, RGB565_BLACK));
        return;
    }

    if (m->count == 0) {
        ESP_ERROR_CHECK(display_draw_text_centered(88, "NO MATCH", RGB565_RED, RGB565_BLACK));
    } else {
        draw_match_row_large(st);
    }

    if (m->count == 1) {
        ESP_ERROR_CHECK(display_draw_text_small(4, 108, "One word left: D0=YES  D2=NO", RGB565_RED, RGB565_BLACK));
    }

    if (st->entered_count < 12) {
        ESP_ERROR_CHECK(display_draw_text_small(4, 118, "D0=letter  D1h=back  D2h=del", RGB565_DIM_GRAY, RGB565_BLACK));
    } else {
        ESP_ERROR_CHECK(display_draw_text_small(4, 118, "12/12: D0 hold=save  D2 hold=del", RGB565_DIM_GRAY, RGB565_BLACK));
    }
}

static int build_phrase_from_entered(ui_state_t *st)
{
    size_t cursor = 0;

    st->phrase[0] = '\0';
    for (size_t i = 0; i < st->entered_count; ++i) {
        size_t wlen = strlen(st->entered_words[i]);

        if (i > 0) {
            if (cursor + 1 >= sizeof(st->phrase)) {
                return -1;
            }
            st->phrase[cursor++] = ' ';
        }

        if (cursor + wlen >= sizeof(st->phrase)) {
            return -1;
        }

        memcpy(st->phrase + cursor, st->entered_words[i], wlen);
        cursor += wlen;
        st->phrase[cursor] = '\0';
    }

    return 0;
}

int ui_run_until_working(char out_mnemonic[256])
{
    ui_state_t st;
    bool main_menu_drawn = false;

    if (out_mnemonic == NULL) {
        return -1;
    }

    memset(&st, 0, sizeof(st));
    st.main_sel = UI_MAIN_GENERATE;

    draw_splash();
    (void)wait_any_button();

    while (1) {
        btn_event_t evt;

        if (!main_menu_drawn) {
            draw_main_menu_full(&st);
            main_menu_drawn = true;
        }

        evt = wait_any_button();

        if (evt.id == BTN_D2) {
            int prev = (int)st.main_sel;
            if (st.has_phrase) {
                st.main_sel = (st.main_sel == UI_MAIN_GENERATE) ? UI_MAIN_VIEW : (ui_main_menu_item_t)((int)st.main_sel - 1);
            } else {
                st.main_sel = (st.main_sel == UI_MAIN_GENERATE) ? UI_MAIN_ENTER : UI_MAIN_GENERATE;
            }
            if ((int)st.main_sel != prev) {
                draw_main_menu_selection_change(st.has_phrase, prev, (int)st.main_sel);
            }
            continue;
        }

        if (evt.id == BTN_D1) {
            int prev = (int)st.main_sel;
            if (st.has_phrase) {
                st.main_sel = (st.main_sel == UI_MAIN_VIEW) ? UI_MAIN_GENERATE : (ui_main_menu_item_t)((int)st.main_sel + 1);
            } else {
                st.main_sel = (st.main_sel == UI_MAIN_ENTER) ? UI_MAIN_GENERATE : UI_MAIN_ENTER;
            }
            if ((int)st.main_sel != prev) {
                draw_main_menu_selection_change(st.has_phrase, prev, (int)st.main_sel);
            }
            continue;
        }

        if (evt.id != BTN_D0) {
            continue;
        }

        if (st.main_sel == UI_MAIN_VIEW) {
            if (!st.has_phrase) {
                continue;
            }

            draw_view_page(&st);
            (void)wait_any_button();
            main_menu_drawn = false;
            continue;
        }

        if (st.main_sel == UI_MAIN_GENERATE) {
            draw_generate_page();
            if (generate_mnemonic_12(st.phrase, sizeof(st.phrase)) != 0) {
                ESP_ERROR_CHECK(display_clear(RGB565_BLACK));
                ESP_ERROR_CHECK(display_draw_text_centered(56, "GEN FAILED", RGB565_RED, RGB565_BLACK));
                vTaskDelay(pdMS_TO_TICKS(2000));
                main_menu_drawn = false;
                continue;
            }

            st.has_phrase = true;
            draw_view_page(&st);
            (void)wait_any_button();

            strncpy(out_mnemonic, st.phrase, 255);
            out_mnemonic[255] = '\0';
            secure_memzero(&st, sizeof(st));
            return 0;
        }

        clear_enter_words(&st);
        st.letter_idx = 0;

        {
            prefix_matches_t matches;
            bool allowed_letters[26];

            compute_prefix_matches(&st, &matches);
            compute_allowed_next_letters(&st, allowed_letters);
            st.letter_idx = first_allowed_letter(allowed_letters);
            draw_enter_compact_static(&st, allowed_letters);
            draw_enter_compact_dynamic(&st, &matches);

            while (1) {
                evt = wait_any_button();

                if (evt.id == BTN_D2 && evt.kind == BTN_EVT_PRESS) {
                    if (matches.count == 1 && st.prefix_len > 0) {
                        st.prefix_len = 0;
                        st.prefix[0] = '\0';
                        compute_prefix_matches(&st, &matches);
                        draw_enter_compact_dynamic(&st, &matches);
                    } else {
                        int prev_idx = st.letter_idx;
                        st.letter_idx = find_next_allowed_letter(st.letter_idx, -1, allowed_letters);
                        draw_enter_letter_cursor_change(prev_idx, st.letter_idx, allowed_letters);
                    }
                    continue;
                }

                if (evt.id == BTN_D1 && evt.kind == BTN_EVT_PRESS) {
                    int prev_idx = st.letter_idx;
                    st.letter_idx = find_next_allowed_letter(st.letter_idx, +1, allowed_letters);
                    draw_enter_letter_cursor_change(prev_idx, st.letter_idx, allowed_letters);
                    continue;
                }

                if (evt.id == BTN_D1 && evt.kind == BTN_EVT_LONG_PRESS) {
                    if (st.prefix_len > 0) {
                        st.prefix_len--;
                        st.prefix[st.prefix_len] = '\0';
                    }
                    compute_prefix_matches(&st, &matches);
                    compute_allowed_next_letters(&st, allowed_letters);
                    st.letter_idx = allowed_letters[st.letter_idx] ? st.letter_idx : first_allowed_letter(allowed_letters);
                    draw_enter_compact_static(&st, allowed_letters);
                    draw_enter_compact_dynamic(&st, &matches);
                    continue;
                }

                if (evt.id == BTN_D2 && evt.kind == BTN_EVT_LONG_PRESS) {
                    if (st.prefix_len > 0) {
                        st.prefix_len = 0;
                        st.prefix[0] = '\0';
                    } else if (st.entered_count > 0) {
                        secure_memzero(st.entered_words[st.entered_count - 1], sizeof(st.entered_words[0]));
                        st.entered_count--;
                    }
                    compute_prefix_matches(&st, &matches);
                    compute_allowed_next_letters(&st, allowed_letters);
                    st.letter_idx = allowed_letters[st.letter_idx] ? st.letter_idx : first_allowed_letter(allowed_letters);
                    draw_enter_compact_static(&st, allowed_letters);
                    draw_enter_compact_dynamic(&st, &matches);
                    continue;
                }

                if (evt.id == BTN_D0 && evt.kind == BTN_EVT_PRESS) {
                    if (st.entered_count >= 12) {
                        if (build_phrase_from_entered(&st) == 0 && mnemonic_validate(st.phrase) == 0) {
                            draw_view_page(&st);
                            (void)wait_any_button();
                            strncpy(out_mnemonic, st.phrase, 255);
                            out_mnemonic[255] = '\0';
                            secure_memzero(&st, sizeof(st));
                            return 0;
                        }
                        continue;
                    }

                    if (matches.count == 1 && st.prefix_len > 0) {
                        const char *w = mnemonic_word_at(matches.first_idx[0]);
                        if (w != NULL) {
                            strncpy(st.entered_words[st.entered_count], w, 11);
                            st.entered_words[st.entered_count][11] = '\0';
                            st.entered_count++;
                        }
                        st.prefix_len = 0;
                        st.prefix[0] = '\0';

                        if (st.entered_count == 12) {
                            if (build_phrase_from_entered(&st) != 0 || mnemonic_validate(st.phrase) != 0) {
                                ESP_ERROR_CHECK(display_clear(RGB565_BLACK));
                                ESP_ERROR_CHECK(display_draw_text_centered(56, "BAD CHECKSUM", RGB565_RED, RGB565_BLACK));
                                vTaskDelay(pdMS_TO_TICKS(1200));
                                compute_prefix_matches(&st, &matches);
                                compute_allowed_next_letters(&st, allowed_letters);
                                st.letter_idx = allowed_letters[st.letter_idx] ? st.letter_idx : first_allowed_letter(allowed_letters);
                                draw_enter_compact_static(&st, allowed_letters);
                                draw_enter_compact_dynamic(&st, &matches);
                                continue;
                            }

                            draw_view_page(&st);
                            (void)wait_any_button();
                            strncpy(out_mnemonic, st.phrase, 255);
                            out_mnemonic[255] = '\0';
                            secure_memzero(&st, sizeof(st));
                            return 0;
                        }

                        compute_prefix_matches(&st, &matches);
                        compute_allowed_next_letters(&st, allowed_letters);
                        st.letter_idx = first_allowed_letter(allowed_letters);
                        draw_enter_compact_static(&st, allowed_letters);
                        draw_enter_compact_dynamic(&st, &matches);
                        continue;
                    }

                    if (st.prefix_len + 1 < sizeof(st.prefix)) {
                        st.prefix[st.prefix_len++] = (char)('a' + st.letter_idx);
                        st.prefix[st.prefix_len] = '\0';

                        compute_prefix_matches(&st, &matches);
                        if (matches.count == 0) {
                            st.prefix_len--;
                            st.prefix[st.prefix_len] = '\0';
                            compute_prefix_matches(&st, &matches);
                        }
                        compute_allowed_next_letters(&st, allowed_letters);
                        st.letter_idx = allowed_letters[st.letter_idx] ? st.letter_idx : first_allowed_letter(allowed_letters);
                        draw_enter_compact_static(&st, allowed_letters);
                        draw_enter_compact_dynamic(&st, &matches);
                    }
                    continue;
                }

                if (evt.id == BTN_D0 && evt.kind == BTN_EVT_LONG_PRESS) {
                    compute_prefix_matches(&st, &matches);

                    if (st.entered_count == 12 && st.prefix_len == 0) {
                        if (build_phrase_from_entered(&st) != 0 || mnemonic_validate(st.phrase) != 0) {
                            ESP_ERROR_CHECK(display_clear(RGB565_BLACK));
                            ESP_ERROR_CHECK(display_draw_text_centered(56, "BAD CHECKSUM", RGB565_RED, RGB565_BLACK));
                            vTaskDelay(pdMS_TO_TICKS(1200));
                            compute_prefix_matches(&st, &matches);
                            compute_allowed_next_letters(&st, allowed_letters);
                            st.letter_idx = allowed_letters[st.letter_idx] ? st.letter_idx : first_allowed_letter(allowed_letters);
                            draw_enter_compact_static(&st, allowed_letters);
                            draw_enter_compact_dynamic(&st, &matches);
                            continue;
                        }

                        strncpy(out_mnemonic, st.phrase, 255);
                        out_mnemonic[255] = '\0';
                        secure_memzero(&st, sizeof(st));
                        return 0;
                    }

                    if (st.prefix_len == 0) {
                        clear_enter_words(&st);
                        break;
                    }

                    continue;
                }
            }
        }

        main_menu_drawn = false;
    }
}