#include "ayla/envelope.hpp"

#include <cstdio>
#include <cstdlib>
#include <cstring>

#include "mbedtls/base64.h"


namespace fgl::ayla {

namespace {

bool b64_encode(const uint8_t* in, size_t in_len, char* out, size_t out_cap) {
  size_t olen = 0;
  if (mbedtls_base64_encode(reinterpret_cast<unsigned char*>(out), out_cap,
                            &olen, in, in_len) != 0) {
    return false;
  }
  return true;  // mbedtls пишет NUL-терминатор
}

bool b64_decode(const char* in, uint8_t* out, size_t out_cap, size_t* out_len) {
  size_t olen = 0;
  if (mbedtls_base64_decode(out, out_cap, &olen,
                            reinterpret_cast<const unsigned char*>(in),
                            strlen(in)) != 0) {
    return false;
  }
  *out_len = olen;
  return true;
}

// Извлекает seq_no верхнего уровня из '{"seq_no":N,"data":...}'.
// Ручной depth-1 сканер (без ограничений на размер data; вход уже
// проверен подписью). Понимает пробелы между лексемами.
bool extract_seq_no(const char* json, int64_t* seq_no) {
  const char* p = json;
  while (*p == ' ' || *p == '\t') p++;
  if (*p != '{') return false;
  p++;
  for (;;) {
    while (*p == ' ' || *p == ',' || *p == '\t') p++;
    if (*p == '}' || *p == '\0') return false;
    // ключ
    if (*p != '"') return false;
    const char* k = ++p;
    while (*p != '"' && *p != '\0') p++;
    if (*p != '"') return false;
    size_t klen = static_cast<size_t>(p - k);
    p++;
    while (*p == ' ' || *p == '\t') p++;
    if (*p != ':') return false;
    p++;
    while (*p == ' ' || *p == '\t') p++;
    if (klen == 6 && strncmp(k, "seq_no", 6) == 0) {
      if (*p == '-' || (*p >= '0' && *p <= '9')) {
        *seq_no = strtoll(p, nullptr, 10);
        return true;
      }
      return false;
    }
    // пропустить значение другого ключа
    if (*p == '{' || *p == '[') {
      int depth = 0;
      bool in_str = false;
      for (; *p != '\0'; p++) {
        if (in_str) {
          if (*p == '\\') {
            p++;
            if (*p == '\0') return false;  // обрыв после escape
          } else if (*p == '"') {
            in_str = false;
          }
          continue;
        }
        if (*p == '"') {
          in_str = true;
        } else if (*p == '{' || *p == '[') {
          depth++;
        } else if (*p == '}' || *p == ']') {
          depth--;
          if (depth == 0) {
            p++;
            break;
          }
        }
      }
      if (*p == '\0') return false;
    } else if (*p == '"') {
      p++;
      while (*p != '"' && *p != '\0') {
        if (*p == '\\') {
          p++;
          if (*p == '\0') return false;  // обрыв после escape
        }
        p++;
      }
      if (*p != '"') return false;
      p++;
    } else {
      while (*p != ',' && *p != '}' && *p != '\0') p++;
      if (*p == '\0') return false;
    }
  }
}

// Сравнение за константное время (аналог mbedtls_ct_memcmp, без зависимости
// от версии библиотеки).
int ct_memcmp(const uint8_t* a, const uint8_t* b, size_t n) {
  uint8_t diff = 0;
  for (size_t i = 0; i < n; i++) diff |= static_cast<uint8_t>(a[i] ^ b[i]);
  return diff != 0 ? 1 : 0;
}

}  // namespace

bool envelope_pack(DirectionCrypto& dir, int64_t seq_no,
                   const char* data_json, char* out, size_t out_cap) {
  if (data_json == nullptr || out == nullptr) return false;

  // Открытый текст: {"seq_no":N,"data":<data_json>}
  char plain[kEnvelopeMaxPlain];
  int n = snprintf(plain, sizeof(plain), "{\"seq_no\":%lld,\"data\":%s}",
                   static_cast<long long>(seq_no), data_json);
  if (n <= 0 || static_cast<size_t>(n) >= sizeof(plain)) return false;
  size_t plain_len = static_cast<size_t>(n);

  // Подпись — по байтам открытого текста без паддинга.
  uint8_t mac[kSha256Len];
  dir.sign(mac, reinterpret_cast<const uint8_t*>(plain), plain_len);
  char sign_b64[64];
  if (!b64_encode(mac, sizeof(mac), sign_b64, sizeof(sign_b64))) return false;

  // Шифрование с zero-паддингом (>=1 NUL, кратно 16).
  uint8_t ct[kEnvelopeMaxPlain + kAesBlockLen];
  long ct_len = dir.encrypt(ct, sizeof(ct),
                            reinterpret_cast<const uint8_t*>(plain), plain_len);
  if (ct_len < 0) return false;
  char enc_b64[kEnvelopeMaxB64];
  if (!b64_encode(ct, static_cast<size_t>(ct_len), enc_b64, sizeof(enc_b64))) {
    return false;
  }

  int m = snprintf(out, out_cap, "{\"enc\":\"%s\",\"sign\":\"%s\"}", enc_b64,
                   sign_b64);
  return m > 0 && static_cast<size_t>(m) < out_cap;
}

bool envelope_unpack(DirectionCrypto& dir, const char* enc_b64,
                     const char* sign_b64, char* json_out, size_t out_cap,
                     int64_t* seq_no) {
  if (enc_b64 == nullptr || sign_b64 == nullptr || json_out == nullptr) {
    return false;
  }
  uint8_t ct[kEnvelopeMaxPlain + kAesBlockLen];
  size_t ct_len = 0;
  if (!b64_decode(enc_b64, ct, sizeof(ct), &ct_len)) return false;

  // Расшифровка: CBC-цепочка двигается независимо от исхода проверки подписи.
  uint8_t plain[kEnvelopeMaxPlain + kAesBlockLen];
  long plain_len = dir.decrypt(plain, sizeof(plain), ct, ct_len);
  if (plain_len < 0 || static_cast<size_t>(plain_len) >= out_cap) return false;

  // Проверка подписи по расшифрованным байтам (без паддинга);
  // сравнение декодированных байт — константное время.
  uint8_t mac[kSha256Len];
  dir.sign(mac, plain, static_cast<size_t>(plain_len));
  uint8_t given[kSha256Len];
  size_t given_len = 0;
  if (!b64_decode(sign_b64, given, sizeof(given), &given_len)) return false;
  if (given_len != kSha256Len) return false;
  if (ct_memcmp(mac, given, kSha256Len) != 0) return false;

  memcpy(json_out, plain, static_cast<size_t>(plain_len));
  json_out[plain_len] = '\0';
  return extract_seq_no(json_out, seq_no);
}

}  // namespace fgl::ayla