#include <LiquidCrystal_I2C.h>
#include <PS2KeyAdvanced.h>
#include <Wire.h> 

#include "enigma_types.h"
#include "HCuOLED.h"
#include "SPI.h"

#define DATAPIN 4
#define IRQPIN  3
#define CS_DI  10
#define DC_DI   9
#define RST_DI  8

PS2KeyAdvanced keyboard;
HCuOLED HCuOLED(SH1106, CS_DI, DC_DI, RST_DI);

LiquidCrystal_I2C lcd(0x27, 16, 2);
String lcd_output;

uint16_t toKeyIndex(const String& input_hex)
{
    uint16_t index = -1;
    for (size_t i = 0; i < ALPHABET_SIZE; ++i)
    {
        if (key_values[i].hex == input_hex)
        {
            index = i;
            break;
        }
    }

    return index;
}

size_t forward1(size_t index, size_t current_wheel)
{
    size_t input = (index + key_shifts[current_wheel]) % ALPHABET_SIZE;

    size_t output = 0;
    for (size_t i = 0; i < ALPHABET_SIZE; ++i)
    {
        if (mutations1[i].from == input)
        {
            output = mutations1[i].to;
            break;
        }
    }

    return output;
}

size_t forward2(size_t index, size_t current_wheel)
{
    size_t input = (index + key_shifts[current_wheel]) % ALPHABET_SIZE;

    size_t output = 0;
    for (size_t i = 0; i < ALPHABET_SIZE; ++i)
    {
        if (mutations2[i].from == input)
        {
            output = mutations2[i].to;
            break;
        }
    }

    return output;
}

size_t forward3(size_t index, size_t current_wheel)
{
    size_t input = (index + key_shifts[current_wheel]) % ALPHABET_SIZE;

    size_t output = 0;
    for (size_t i = 0; i < ALPHABET_SIZE; ++i)
    {
        if (mutations3[i].from == input)
        {
            output = mutations3[i].to;
            break;
        }
    }

    return output;
}

size_t forward4(size_t index)
{
    size_t input = index % ALPHABET_SIZE;

    size_t output = 0;
    for (size_t i = 0; i < ALPHABET_SIZE; ++i)
    {
        if (mutations4[i].from == input)
        {
            output = mutations4[i].to;
            break;
        }
    }

    return output;
}

size_t backward1(size_t index, size_t current_wheel)
{
    int output = 0;

    for (size_t i = 0; i < ALPHABET_SIZE; ++i)
    {
        if (index == mutations1[i].to)
        {
            output = (mutations1[i].from - key_shifts[current_wheel]);
            while (output < 0)
            {
                output += ALPHABET_SIZE;
            }

            output = output % ALPHABET_SIZE;
        }
    }

    return output;
}

size_t backward2(size_t index, size_t current_wheel)
{
    int output = 0;

    for (size_t i = 0; i < ALPHABET_SIZE; ++i)
    {
        if (index == mutations2[i].to)
        {
            output = (mutations2[i].from - key_shifts[current_wheel]);
            while (output < 0)
            {
                output += ALPHABET_SIZE;
            }

            output = output % ALPHABET_SIZE;
        }
    }

    return output;
}

size_t backward3(size_t index, size_t current_wheel)
{
    int output = 0;

    for (size_t i = 0; i < ALPHABET_SIZE; ++i)
    {
        if (index == mutations3[i].to)
        {
            output = (mutations3[i].from - key_shifts[current_wheel]);
            while (output < 0)
            {
                output += ALPHABET_SIZE;
            }

            output = output % ALPHABET_SIZE;
        }
    }

    return output;
}

size_t encode(size_t index)
{
    index = forward1(index, 0);
    index = forward2(index, 1);
    index = forward3(index, 2);
    index = forward4(index);
    index = backward3(index, 2);
    index = backward2(index, 1);
    index = backward1(index, 0);

    return index;
}


void rotate()
{
    ++key_shifts[0];
    if (key_shifts[0] == 26)
    {
        key_shifts[0] = 0;
        ++key_shifts[1];
        if (key_shifts[1] == 26)
        {
            key_shifts[1] = 0;
            ++key_shifts[2];
            if (key_shifts[2] == 26)
            {
                key_shifts[2] = 0;
            }
        }
    }
}

void setup()
{
    HCuOLED.Reset();
    keyboard.begin(DATAPIN, IRQPIN);
    Serial.begin(115200);

    lcd.begin();
    lcd.backlight();
}

void loop()
{
    if (!keyboard.available())
        return;

    const uint16_t key = keyboard.read();
    if (key <= 0)
        return;

    const String value = String(key, HEX);
    const uint16_t index = toKeyIndex(value);
    if (index == -1)
        return;

    size_t encoded_index = encode(index);
    const String encoded_letter = key_values[encoded_index].view;
    
    HCuOLED.SetFont(MedProp_11pt);
    HCuOLED.Cursor(5, 5);
    HCuOLED.Print(key_shifts[0]);
    Serial.print(key_shifts[0]);

    HCuOLED.Cursor(5, 25);
    HCuOLED.Print(key_shifts[1]);
    Serial.print(key_shifts[1]);

    HCuOLED.Cursor(5, 40);
    HCuOLED.Print(key_shifts[2]);
    Serial.print(key_shifts[2]);
    HCuOLED.Refresh();
    HCuOLED.ClearBuffer();
    Serial.println(" ");

    if (lcd_output.length() == 16)
    {
        lcd.setCursor(0, 0);
        lcd_output = "";
    }

    lcd_output = (lcd_output + encoded_letter);

    lcd.clear();
    lcd.print(lcd_output);
    rotate();
}