#pragma once

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

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

namespace Blowfish {

template <std::size_t sz>
concept is_valid_key_length = (KEYLEN_MIN <= sz && sz <= KEYLEN_MAX);

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

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

template <std::size_t keylen>
requires is_valid_key_length<keylen> class Context {
public:
  constexpr Context(const std::span<const uint8_t> &key) {
    for (std::size_t i = 0; i < BOXES; ++i) {
      std::copy(S_INIT[i].begin(), S_INIT[i].end(), S[i].begin());
    }

    for (std::size_t i = 0; i < SUBKEYS; ++i) {
        const auto idx0 = (i * 4 + 0) % keylen;
        const auto idx1 = (i * 4 + 1) % keylen;
        const auto idx2 = (i * 4 + 2) % keylen;
        const auto idx3 = (i * 4 + 3) % keylen;
        const uint32_t qk = (static_cast<uint32_t>(key[idx0]) << 24) +
                            (static_cast<uint32_t>(key[idx1]) << 16) +
                            (static_cast<uint32_t>(key[idx2]) << 8) +
                            static_cast<uint32_t>(key[idx3]);
        P[i] = P_INIT[i] ^ qk;
    }

    Block x(0, 0);

    for (std::size_t i = 0; i < SUBKEYS; i += 2) {
      x = encrypt(x);
      P[i] = x.L;
      P[i + 1] = x.R;
    }

    for (std::size_t i = 0; i < BOXES; ++i) {
      for (std::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 (std::size_t i = 0; i < ROUNDS; ++i)
      [[likely]] {
        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 (std::size_t i = ROUNDS + 1; i > 1; --i)
      [[likely]] {
        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;
  }

  std::array<uint32_t, SUBKEYS> P;
  std::array<std::array<uint32_t, ENTRIES>, BOXES> S;
};

} // namespace Blowfish