#pragma once

// https://www.schneier.com/academic/archives/1994/09/description_of_a_new.html

#include "blowfish_init.hpp"
#include <algorithm>

class Blowfish {
public:
  struct Block {
    constexpr Block(uint64_t x) {
      L = x >> 32;
      R = x & 0xffffffff;
    }
    constexpr Block(uint32_t l, uint32_t r) : L(l), R(r) {}

    constexpr operator uint64_t() const {
      return (static_cast<uint64_t>(L) << 32) + R;
    }
    uint32_t L;
    uint32_t R;
  };

  constexpr Blowfish(const uint8_t key[], size_t keylen) {
    for (size_t i = 0; i < BOXES; ++i) {
      for (size_t j = 0; j < ENTRIES; ++j) {
        S[i][j] = S_INIT[i][j];
      }
    }

    //
    size_t k = 0;
    for (size_t i = 0; i < SUBKEYS; ++i) {
      uint32_t qk = 0;
      for (size_t j = 0; j < 4; ++j) {
        qk = (qk << 8) | key[k];
        ++k;
        if (k >= keylen) {
          k = 0;
        }
      }
      P[i] = P_INIT[i] ^ qk;
    }

    Block x(0, 0);

    for (size_t i = 0; i < SUBKEYS; i += 2) {
      x = encrypt(x);
      P[i] = x.L;
      P[i + 1] = x.R;
    }
    for (size_t i = 0; i < BOXES; ++i) {
      for (size_t j = 0; j < ENTRIES; j += 2) {
        x = encrypt(x);
        S[i][j] = x.L;
        S[i][j + 1] = x.R;
      }
    }
  }

  constexpr Block encrypt(const Block &x) const {
    Block y(x);
    for (size_t i = 0; i < ROUNDS; ++i) {
      y.L = y.L ^ P[i];
      y.R = F(y.L) ^ y.R;
      std::swap(y.L, y.R);
    }
    std::swap(y.L, y.R);
    y.R = y.R ^ P[16];
    y.L = y.L ^ P[17];

    return y;
  }

  constexpr Block decrypt(const Block &x) const {
    Block y(x);
    for (size_t i = ROUNDS + 1; i > 1; --i) {
      y.L = y.L ^ P[i];
      y.R = F(y.L) ^ y.R;
      std::swap(y.L, y.R);
    }
    std::swap(y.L, y.R);
    y.R = y.R ^ P[1];
    y.L = y.L ^ P[0];
    return y;
  }

private:
  constexpr uint32_t F(uint32_t x) const {
    const uint8_t a = (x >> 24) & 0xff;
    const uint8_t b = (x >> 16) & 0xff;
    const uint8_t c = (x >> 8) & 0xff;
    const uint8_t d = x & 0xff;

    uint32_t y = S[0][a] + S[1][b];
    y = y ^ S[2][c];
    y = y + S[3][d];
    return y;
  }

  uint32_t P[SUBKEYS];
  uint32_t S[BOXES][ENTRIES];
};