#include <Arduino.h>

#include <Adafruit_NeoPixel.h>

#include <MFRC522v2.h>
#include <MFRC522DriverSPI.h>
#include <MFRC522DriverPinSimple.h>
#include <MFRC522Hack.h>

Adafruit_NeoPixel pixel(1, 12, NEO_GRB + NEO_KHZ800);

static MFRC522DriverPinSimple cs_pin(D0);
static MFRC522DriverSPI driver(cs_pin);
static MFRC522 mfrc522(driver);
static MFRC522Hack hack(mfrc522, true, &Serial);

static MFRC522::MIFARE_Key key;

enum class Mode {
    Read,
    Write,
    Unbrick,
};

constexpr uint8_t BTN_1 = D6, BTN_2 = D5, BTN_3 = D4;
constexpr uint8_t IRQ = D7;

static String display_mode(const Mode& mode) {
  switch (mode) {
    case Mode::Read:
      return "read";
    
    case Mode::Write:
      return "write";
    
    case Mode::Unbrick:
      return "unbrick";
  }

  return "<unknown>";
}

constexpr uint8_t DEFAULT_KEY[4] = {0, 0, 0, 0};
static uint8_t key_bytes[4];

constexpr uint8_t OVERWRITE_KEY[4] = {0xa5, 0xa5, 0xa5, 0xa5};

static bool keyIsDefault() {
  return !memcmp(key_bytes, DEFAULT_KEY, sizeof(DEFAULT_KEY));
}

static bool keyIsOverwrite() {
  return !memcmp(key_bytes, OVERWRITE_KEY, sizeof(OVERWRITE_KEY));
}

static std::vector<char> BUF(128);
static Mode mode = Mode::Read;

bool parseSetKey(const String& s);
static void printSelectedUid();

constexpr float LED_SCALE = 0.3;

static uint8_t R, G, B;

struct Color {
  uint8_t r, g, b;
};

constexpr Color
  READ_NODATA = Color { 0, 255, 255 },
  READ = Color { 0, 0, 255 },
  WRITE = Color { 255, 0, 255 },
  WRITE_ERASE = Color { 255, 32, 0 },
  UNBRICK = Color { 255, 255, 0 },
  ERROR = Color { 255, 0, 0 };

constexpr uint8_t rescale_color(uint8_t channel) {
  return (uint8_t)((float)channel * LED_SCALE);
}

static void setColor(uint8_t r, uint8_t g, uint8_t b, bool retain = true) {
    pixel.clear();
    pixel.setPixelColor(0, pixel.Color(rescale_color(r), rescale_color(g), rescale_color(b)));
    pixel.show();

    if (retain) {
      R = r;
      G = g;
      B = b;
    }
}

static void setColor(const Color& color, bool retain = true) {
  setColor(color.r, color.g, color.b, retain);
}

static void restoreColor() {
  setColor(R, G, B);
}

static void updateColor(const Mode& mode) {
  switch (mode) {
      case Mode::Read:
          if (keyIsDefault() || keyIsOverwrite()) setColor(READ_NODATA);
          else setColor(READ);
        
          break;
      
      case Mode::Write:
          if (keyIsOverwrite()) setColor(WRITE_ERASE);
          else setColor(WRITE);
          break;

      case Mode::Unbrick:
          setColor(UNBRICK);
          break;
  }
}

void handleInterrupt() {
  auto read = digitalRead(BTN_1) == LOW;
  auto write = digitalRead(BTN_2) == LOW;
  auto unbrick = digitalRead(BTN_3) == LOW;

  if (read && write) {
    memset(key_bytes, 0, sizeof(key_bytes));
    Serial.println("[int] erased key");
    updateColor(mode);
    return;
  }

  if (write && unbrick) {
    memset(key_bytes, 0xa5, sizeof(key_bytes));
    Serial.println("[int] entered write-zero mode");
    mode = Mode::Write;
    updateColor(mode);
    return;
  }

  if (write) {
    if (keyIsDefault()) {
      Serial.println("unable to enter write mode: null key");

      for (int i = 0; i < 5; i++)  {
        setColor(ERROR, false);
        delay(20);

        setColor(0, 0, 0, false);
        delay(20);
      }

      restoreColor();

      return;
    }

    mode = Mode::Write;
    updateColor(mode);
    Serial.println("[int] mode: write");
    return;
  }

  if (read) {
    mode = Mode::Read;
    updateColor(mode);
    Serial.println("[int] mode: read");
    return;
  }

  if (unbrick) {
    mode = Mode::Unbrick;
    updateColor(mode);
    Serial.println("[int] mode: unbrick");
  }
}


void setup() {
    Serial.begin(115200);

    pixel.begin();

    pinMode(11, OUTPUT);
    digitalWrite(11, HIGH);

    pinMode(BTN_1, INPUT_PULLUP);
    pinMode(BTN_2, INPUT_PULLUP);
    pinMode(BTN_3, INPUT_PULLUP);
    pinMode(IRQ, INPUT_PULLDOWN);

    attachInterrupt(digitalPinToInterrupt(BTN_1), handleInterrupt, FALLING);
    attachInterrupt(digitalPinToInterrupt(BTN_2), handleInterrupt, FALLING);
    attachInterrupt(digitalPinToInterrupt(BTN_3), handleInterrupt, FALLING);

    attachInterrupt(digitalPinToInterrupt(IRQ), []() {
      Serial.println("IRQ");
    }, CHANGE);
    
    setColor(READ_NODATA);

    mfrc522.PCD_Init();

    memset(key.keyByte, 0xff, sizeof(key.keyByte));
    memset(key_bytes, 0x0, sizeof(key_bytes));

    BUF.clear();

    Serial.println("\n\n\n");
    Serial.println("card cloner start ok");
    Serial.println("mode: read");
}

void loop() {
    unsigned long delayTime = 15;
    bool status;

    auto cardPresent = mfrc522.PICC_IsNewCardPresent();
    if (!cardPresent) {
      status = true;
      goto end;
    }

    pixel.clear();
    pixel.show();

    switch (mode) {
        case Mode::Read:
            if (!mfrc522.PICC_ReadCardSerial()) goto end;

            Serial.printf("read card (sak: 0x%02x), uid: ", mfrc522.uid.sak);
            printSelectedUid();
            Serial.println();

            memcpy(key_bytes, mfrc522.uid.uidByte, MIN(mfrc522.uid.size, sizeof(key_bytes)));
            updateColor(mode);
            setColor(0, 0, 0, false);

            delayTime = 300;
            status = true;

            break;

        case Mode::Write:
            if (!mfrc522.PICC_ReadCardSerial()) goto end;

            Serial.print("overwriting card: ");
            printSelectedUid();
            Serial.println();

            delayTime = 500;

            status = hack.MIFARE_SetUid(key_bytes, sizeof(key_bytes), key, true);
            mfrc522.PICC_HaltA();
            
            if (!status) {
              Serial.println("failed to write card");
              goto end;
            }

            Serial.println("wrote card");

            if (!mfrc522.PICC_IsNewCardPresent() || !mfrc522.PICC_ReadCardSerial()) {
              Serial.println("unable to readback card serial");
              status = false;
              goto end;
            }

            Serial.print("new serial: ");
            printSelectedUid();
            Serial.println();

            status = true;

            break;

        case Mode::Unbrick:
            delayTime = 500;

            status = hack.MIFARE_UnbrickUidSector();

            if (status) Serial.println(F("unbricked uid sector"));
            else {
                Serial.println(F("failed to unbrick uid sector"));
                mfrc522.PICC_HaltA();
            }

            break;
    }

    end:
        if (!status) setColor(ERROR, false);

        delay(delayTime / 2);
        restoreColor();
        delay(delayTime / 2);
}

static void printSelectedUid() {
    for (int i = 0; i < mfrc522.uid.size; i++) Serial.printf("%02x", mfrc522.uid.uidByte[i]);
}

void serialEvent() {
  while (Serial.available()) {
    auto c = (char)Serial.read();
    auto new_mode = mode;

    if (c != '\n') {
        BUF.push_back(c);
        continue;
    }

    auto s = String(BUF.data(), BUF.size());
    BUF.clear();

    auto defaultKey = keyIsDefault();

    if (s == "read") {
        new_mode = Mode::Read;
    } else if (s == "write") {
        if (defaultKey) {
            Serial.println("unable to enter write mode with null key");
            continue;
        }

        new_mode = Mode::Write;
    } else if (s == "unbrick") {
        new_mode = Mode::Unbrick;
    } else if (s.startsWith("set ")) parseSetKey(s);
    else if (s == "get") {
        for (int i = 0; i < sizeof(key_bytes); i++) Serial.printf("%02x", key_bytes[i]);
        Serial.println();
    } else {
        Serial.printf("unrecognized command: '%s'\n", s);
    }

    if (new_mode == mode) continue;

    Serial.print("mode: ");
    Serial.print(display_mode(new_mode));
    Serial.println();

    mode = new_mode;

    updateColor(mode);
  }
}

bool parseSetKey(const String& s) {
    auto id = s.substring(String("set ").length());
    id.trim();
    id.replace(" ", "");

    if (id.length() != 4 * 2) {
        Serial.println("invalid key (expected 4 bytes as hex, no 0x prefix)");
        return false;
    }

    Serial.print("id: '");
    Serial.print(id);
    Serial.println("'");

    Serial.print("set key: ");

    for (int i = 0; i < 4; i++) {
        auto p = &id.c_str()[i * 2];
        char dat[3] = {p[0], p[1], 0};

        auto result = strtoul(dat, NULL, 16);
        key_bytes[i] = (uint8_t)result;

        Serial.printf("%02x", result);
    }

    Serial.println();

    return true;
}