#include "step4.h"

#define ONE_SECOND_TIME 90'000
#define THREE_SECOND_TIME 90'000
#define SIX_SECOND_TIME 75'000

static const char *TAG = "step4";

static int tone = 0;
static int times = 0;

static const char* TONE_FILES[] = {
    MOUNT_POINT "/low-1.pcm",
    MOUNT_POINT "/low-3.pcm",
    MOUNT_POINT "/low-6.pcm",
    MOUNT_POINT "/high-1.pcm",
    MOUNT_POINT "/high-3.pcm",
    MOUNT_POINT "/high-6.pcm",
};

static const char* LCD_STRINGS[] = {
    "something",
    "nothing",
    "",
    "a word",
    "somethink",
    "what?",
    "LCD",
    "display",
};

static int indicator_led_idxs[LED_COUNT] = {0};

static bool contains_coconut = false;
static char* COCONUT = "coconut";
static char lcd_random_char_set[] = "aeiou tnsrhldm";
static char random_lcd_text[21] = {0};

static std::random_device rd;
static std::mt19937 gen(rd());
static std::uniform_int_distribution<> tone_dist(0, 5);
static std::uniform_int_distribution<> color_dist(0, 6);
static std::uniform_int_distribution<> lcd_string_dist(0, 7);
static std::uniform_int_distribution<> lcd_number_dist(0, 15);
static std::uniform_int_distribution<> lcd_rand_char_dist(0, sizeof(lcd_random_char_set)-2);
static std::uniform_int_distribution<> has_coconut_dist(0, 2);
static std::uniform_int_distribution<> coconut_position_dist(0, 13);

static uint8_t NEOPIXEL_COLORS[7][3] = {
    {20, 0, 0},  // red
    {20, 10, 0}, // orange
    {20, 20, 0}, // yellow
    {0, 20, 0},  // green
    {0, 0, 20},  // blue
    {20, 0, 20}, // purple
    {0, 0, 0}, // off
};

static bool one_second();
static bool three_second();
static bool six_second();

void step4(void) {
    StarCodeHandler star_codes[] = {
        {
            .code = "*1642",
            .display_text = "Starting...",
            .should_exit = true,
            .callback = nullptr,
        },
    };
    int len = sizeof(star_codes)/sizeof(StarCodeHandler);
    do_star_codes(star_codes, len);

    while (times < 4) {
        tone = tone_dist(gen);
        // tone = 2;
        while (get_pressed_button(nullptr)) vTaskDelay(pdMS_TO_TICKS(10));

        play_raw(MOUNT_POINT "/que.pcm");
        play_raw(TONE_FILES[tone]);
        
        bool correct = false;
        switch (tone % 3) {
            case 0:
            correct = one_second();
            break;
            case 1:
            correct = three_second();
            break;
            case 2:
            correct = six_second();
            break;
        }
        if (correct) {
            times++;
            clean_bomb();
            play_raw(MOUNT_POINT "/correct.pcm");            
        } else {
            vTaskDelay(pdMS_TO_TICKS(1500));
        }
        vTaskDelay(pdMS_TO_TICKS(3000));
    }
}

static void generate_random_lcd_text(void) {
    for (int i = 0; i < 20; i++) {
        int char_idx = lcd_rand_char_dist(gen);
        random_lcd_text[i] = lcd_random_char_set[char_idx];
    }
    contains_coconut = (has_coconut_dist(gen) == 0);
    if (contains_coconut) {
        int idx = coconut_position_dist(gen);
        for (int i = 0; i < 7; i++) {
            random_lcd_text[idx+i] = COCONUT[i];
            // ESP_LOGI(TAG, "Writing idx %d to %c. Is %c", idx+i, COCONUT[i], random_lcd_text[idx+i]);
        }
        // ESP_LOGI(TAG, "Now: %s", random_lcd_text);
    }
}

/// Sets the leds to random values.
///
/// This does not flush the leds.
static void rng_leds() {
    for (int i = 0; i < LED_COUNT; i++) {
        indicator_led_idxs[i] = color_dist(gen);
    }
}

static void write_leds() {
    // update all the leds
    for (int i = 0; i < LED_COUNT; i++) {
        auto colors = NEOPIXEL_COLORS[indicator_led_idxs[i]];
        led_strip_set_pixel(leds, i, colors[0], colors[1], colors[2]);
    }
    led_strip_refresh(leds);
}

static uint8_t four_bit_flag(bool b0, bool b1, bool b2, bool b3) {
    return
        (b0 << 0) |
        (b1 << 1) |
        (b2 << 2) |
        (b3 << 3)
    ;
}

static void print_bin(char* out_str, uint8_t n) {
    out_str[0] = ((n & 0b1000) ? '1' : '0');
    out_str[1] = ((n & 0b0100) ? '1' : '0');
    out_str[2] = ((n & 0b0010) ? '1' : '0');
    out_str[3] = ((n & 0b0001) ? '1' : '0');
    out_str[4] = ' ';
    out_str[5] = 'i';
    out_str[6] = 'n';
    out_str[7] = ' ';
    out_str[8] = 'o';
    out_str[9] = 'r';
    out_str[10] = 'd';
    out_str[11] = 'e';
    out_str[12] = 'r';
    out_str[13] = ':';
    out_str[14] = ' ';
    out_str[15] = ((n & 0b0001) ? '1' : '0');
    out_str[16] = ((n & 0b0010) ? '1' : '0');
    out_str[17] = ((n & 0b0100) ? '1' : '0');
    out_str[18] = ((n & 0b1000) ? '1' : '0');
}

static void debug_correct_values(uint8_t correct_buttons, uint8_t button_mask, uint8_t correct_switches) {
    char buf[20] = {0};
    print_bin(buf, correct_switches);
    ESP_LOGI(TAG, "Expected Switch State: 0b%s", buf);
    print_bin(buf, correct_buttons);
    ESP_LOGI(TAG, "Expected Button State: 0b%s", buf);
    print_bin(buf, button_mask);
    ESP_LOGI(TAG, "Button Mask:           0b%s", buf);
}

static void debug_actual_values(uint8_t buttons, uint8_t switch_) {
    char buf[20] = {0};
    print_bin(buf, switch_);
    ESP_LOGI(TAG, "Actual Switch State:  0b%s", buf);
    print_bin(buf, buttons);
    ESP_LOGI(TAG, "Actual Button State:  0b%s", buf);
    ESP_LOGI(TAG, "");
}

static void wait_for_timer(void) {
    KeypadKey key;
    while (get_module_time() > 0) {
        if (get_pressed_keypad(&key) && key == KeypadKey::kd) {
            set_module_time(0);
            return;
        }
        vTaskDelay(pdMS_TO_TICKS(100));
    }
}

static bool one_second() {
    clean_bomb();
    set_module_time(ONE_SECOND_TIME);
    start_module_timer();

    rng_leds();
    int speaker_color = indicator_led_idxs[Led::speaker];
    int lcd_string_idx = lcd_string_dist(gen);
    bool was_high = (tone / 3) == 1;
    write_leds();
    lcd_clear(&lcd);
    lcd_set_cursor(&lcd, 1, 1);
    lcd_print(&lcd, LCD_STRINGS[lcd_string_idx]);

    int red_led_count = 0;
    int blue_led_count = 0;
    for (int i = 0; i < LED_COUNT; i++) {
        if (indicator_led_idxs[i] == 0) {
            red_led_count++;
        } else if (indicator_led_idxs[i] == 4) {
            blue_led_count++;
        }
    }

    uint8_t correct_switches = four_bit_flag(
        speaker_color == 0 || speaker_color == 1 || speaker_color == 2,
        lcd_string_idx == 0 || lcd_string_idx == 1,
        was_high,
        !was_high
    );

    uint8_t correct_button_mask = 0b1011;
    uint8_t correct_buttons = four_bit_flag(
        indicator_led_idxs[Led::char_lcd] != 6, // green
        red_led_count > blue_led_count, // red
        0, // yellow UNCHECKED
        indicator_led_idxs[Led::rfid] == 4 || indicator_led_idxs[Led::rfid] == 6 // blue
    );

    debug_correct_values(correct_buttons, correct_button_mask, correct_switches);

    wait_for_timer();

    debug_actual_values(get_button_state(), get_switch_state());

    if (get_switch_state() != correct_switches) {
        clean_bomb();
        strike("Incorrect Switches");
        return false;
    } 
    if ((get_button_state() & correct_button_mask) != correct_buttons) {
        clean_bomb();
        strike("Incorrect Buttons");
        return false;
    }

    return true;
}

static bool three_second() {
    clean_bomb();
    set_module_time(THREE_SECOND_TIME);
    start_module_timer();

    int lcd_number = lcd_number_dist(gen);
    char lcd_number_string[9] = {0};
    sprintf(lcd_number_string, "%d", lcd_number);
    lcd_set_cursor(&lcd, 1, 1);
    lcd_print(&lcd, lcd_number_string);

    bool was_high = (tone / 3) == 1;

    rng_leds();
    write_leds();
    int red_led_count = 0;
    int blue_led_count = 0;
    for (int i = 0; i < LED_COUNT; i++) {
        if (indicator_led_idxs[i] == 0) {
            red_led_count++;
        } else if (indicator_led_idxs[i] == 4) {
            blue_led_count++;
        }
    }

    // reverse the ordering of the bits
    uint8_t correct_switches = four_bit_flag(
        (lcd_number >> 3) & 1,
        (lcd_number >> 2) & 1,
        (lcd_number >> 1) & 1,
        (lcd_number >> 0) & 1
    );
    if (!was_high) {
        correct_switches = (~correct_switches) & 0b1111;
    }

    uint8_t correct_button_mask = 0b1110;
    uint8_t correct_buttons = four_bit_flag(
        0, // green UNCHECKED
        was_high, // red
        (lcd_number % 2) == 0, // yellow
        blue_led_count > red_led_count // blue
    );

    debug_correct_values(correct_buttons, correct_button_mask, correct_switches);

    wait_for_timer();

    debug_actual_values(get_button_state(), get_switch_state());

    if (get_switch_state() != correct_switches) {
        clean_bomb();
        strike("Incorrect Switches");
        return false;
    } 
    if ((get_button_state() & correct_button_mask) != correct_buttons) {
        clean_bomb();
        strike("Incorrect Buttons");
        return false;
    }

    return true;
}

static bool six_second() {
    clean_bomb();
    set_module_time(SIX_SECOND_TIME);
    start_module_timer();

    generate_random_lcd_text();
    vTaskDelay(pdMS_TO_TICKS(10));
    lcd_set_cursor(&lcd, 0, 0);
    lcd_print(&lcd, random_lcd_text);

    int vowels = 0;
    for (int i = 0; i < 20; i++) {
        char c = random_lcd_text[i];
        if (c == 'a' || c == 'e' || c == 'i' || c == 'o' || c == 'u') {
            vowels++;
        }
    }

    bool was_high = (tone / 3) == 1;

    bool second_switch_correct_state = (indicator_led_idxs[Led::switch2] == 0) || (indicator_led_idxs[Led::switch2] == 6);
    second_switch_correct_state = second_switch_correct_state || was_high;

    rng_leds();
    write_leds();

    int green_led_count = 0;
    int blue_led_count = 0;
    for (int i = 0; i < LED_COUNT; i++) {
        if (indicator_led_idxs[i] == 4) {
            blue_led_count++;
        } else if (indicator_led_idxs[i] == 3) {
            green_led_count++;
        }
    }

    int purple_led_on_bottom_count = 0;
    for (int i = Led::rfid; i < LED_COUNT; i++) {
        if (indicator_led_idxs[i] == 5) {
            purple_led_on_bottom_count++;
        }
    }

    uint8_t correct_switches = four_bit_flag(
        vowels > 7,
        second_switch_correct_state,
        true,
        !(purple_led_on_bottom_count > 1)
    );

    uint8_t correct_button_mask = 0b1101;
    uint8_t correct_buttons = four_bit_flag(
        (!was_high) || (green_led_count >= 2) || indicator_led_idxs[Led::keypad] == 4, // green
        0, // red UNCHECKED
        blue_led_count >= 3, // yellow
        contains_coconut // blue
    );

    debug_correct_values(correct_buttons, correct_button_mask, correct_switches);

    wait_for_timer();
    
    debug_actual_values(get_button_state(), get_switch_state());

    if (get_switch_state() != correct_switches) {
        clean_bomb();
        strike("Incorrect Switches");
        return false;
    } 
    if ((get_button_state() & correct_button_mask) != correct_buttons) {
        clean_bomb();
        strike("Incorrect Buttons");
        return false;
    }

    return true;
}