#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 forward(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 (mutations[current_wheel][i].from == input)
        {
            output = mutations[current_wheel][i].to;
            break;
        }
    }

    return output;
}

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

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

            output = output % ALPHABET_SIZE;
        }
    }

    return output;
}

size_t encode(size_t index)
{
    for (size_t i = 0; i < WHEELS_AMOUNT; ++i)
    {
        index = forward(index, i);
    }

    size_t overturning_input = index % ALPHABET_SIZE;
    for (size_t i = 0; i < ALPHABET_SIZE; ++i)
    {
        if (mutations[3][i].from == overturning_input)
        {
            index = mutations[3][i].to;
            break;
        }
    }

    for (size_t i = (WHEELS_AMOUNT - 1); i >= 0; --i)
    {
        index = backward(index, i);
    }

    return index;
}

void rotate()
{
    for (size_t i = 0; i < WHEELS_AMOUNT; ++i)
    {
        ++key_shifts[i];
        if (key_shifts[i] != ALPHABET_SIZE)
            key_shifts[i] = 0;
        else
            break;
    }
}

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;

    size_t y = 0;
    for (size_t i = 0; i < WHEELS_AMOUNT; ++i)
    {
        HCuOLED.Cursor(4, y);
        HCuOLED.SetFont(MedProp_11pt);
        HCuOLED.Print(key_shifts[i]);
        y += 20;
    }

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

    lcd_output = (lcd_output + encoded_letter);

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