#include "jwt.h"

#include <openssl/bio.h>
#include <openssl/buffer.h>
#include <openssl/evp.h>
#include <openssl/hmac.h>
#include <string.h>

// Удаляет символы '=' в конце строки (падинг)
static void remove_padding(const char *str) {
  char *pos = strchr(str, '\0');
  while (pos > str && *(pos - 1) == '=') {
    *(--pos) = '\0';
  }
}

// Преобразует Base64 в Base64URL
static void base64_to_base64url(char *str) {
  for (char *p = str; *p; p++) {
    if (*p == '+')
      *p = '-';
    else if (*p == '/')
      *p = '_';
  }
  remove_padding(str);  // Удаляем паддинг
}

// Преобразует Base64URL в Base64
static void base64url_to_base64(char *str) {
  for (char *p = str; *p; p++) {
    if (*p == '-')
      *p = '+';
    else if (*p == '_')
      *p = '/';
  }
  // Добавляем паддинг для кратности 4
  size_t len = strlen(str);
  size_t pad = (4 - (len % 4)) % 4;
  if (pad) {
    char *end = str + len;
    for (size_t i = 0; i < pad; i++) {
      *(end++) = '=';
    }
    *end = '\0';
  }
}

// Base64URL кодирование
char *jwt_base64url_encode(const unsigned char *input, const int length) {
  BIO *b64 = BIO_new(BIO_f_base64());
  BIO *bio = BIO_new(BIO_s_mem());
  bio = BIO_push(b64, bio);

  BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);
  BIO_write(bio, input, length);
  BIO_flush(bio);

  BUF_MEM *bufferPtr;
  BIO_get_mem_ptr(bio, &bufferPtr);

  const auto result = (char *)malloc(bufferPtr->length + 1);
  memcpy(result, bufferPtr->data, bufferPtr->length);
  result[bufferPtr->length] = '\0';

  BIO_free_all(bio);

  base64_to_base64url(result);  // Преобразование в Base64URL
  return result;
}

// Base64URL декодирование
unsigned char *jwt_base64url_decode(const char *input, int input_len,
                                    int *output_len) {
  if (input_len == 0) input_len = strlen(input);

  // Создаем копию для преобразования
  const auto input_copy = (char *)malloc(input_len + 1);
  strncpy(input_copy, input, input_len);
  input_copy[input_len] = '\0';

  base64url_to_base64(input_copy);  // Преобразование в Base64

  BIO *b64 = BIO_new(BIO_f_base64());
  BIO *bio = BIO_new_mem_buf(input_copy, -1);
  bio = BIO_push(b64, bio);

  BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);

  const auto buffer = (unsigned char *)malloc(input_len);
  const int len = BIO_read(bio, buffer, input_len);

  BIO_free_all(bio);
  free(input_copy);

  if (len <= 0) {
    free(buffer);
    *output_len = 0;
    return nullptr;
  }

  *output_len = len;
  return buffer;
}

// Создание JWT токена
char *jwt_createJWT(const char *header_json, const char *payload_json,
                    const char *secret) {
  char *encodedHeader = jwt_base64url_encode((const unsigned char *)header_json,
                                             strlen(header_json));
  char *encodedPayload = jwt_base64url_encode(
      (const unsigned char *)payload_json, strlen(payload_json));

  char *unsignedToken =
      (char *)malloc(strlen(encodedHeader) + strlen(encodedPayload) + 2);
  sprintf(unsignedToken, "%s.%s", encodedHeader, encodedPayload);

  unsigned int digest_len;
  unsigned char *digest = jwt_hmac_sha256(secret, strlen(secret), unsignedToken,
                                          strlen(unsignedToken), &digest_len);
  char *encodedSignature = jwt_base64url_encode(digest, digest_len);

  char *jwt =
      (char *)malloc(strlen(unsignedToken) + strlen(encodedSignature) + 2);
  sprintf(jwt, "%s.%s", unsignedToken, encodedSignature);

  free(encodedHeader);
  free(encodedPayload);
  free(unsignedToken);
  free(digest);
  free(encodedSignature);

  return jwt;
}

// Генерация HMAC-SHA256
unsigned char *jwt_hmac_sha256(const char *key, int key_len, const char *data,
                               int data_len, unsigned int *digest_len) {
  unsigned char *digest = (unsigned char *)malloc(EVP_MAX_MD_SIZE);
  HMAC(EVP_sha256(), key, key_len, (const unsigned char *)data, data_len,
       digest, digest_len);
  return digest;
}

// Верификация JWT токена
bool jwt_verifyJWT(const char *token, const char *secret) {
  const char *dot1 = strchr(token, '.');
  const char *dot2 = dot1 ? strchr(dot1 + 1, '.') : nullptr;
  if (!dot1 || !dot2) return false;

  size_t header_len = dot1 - token;
  size_t payload_len = dot2 - dot1 - 1;
  size_t sig_len = strlen(dot2 + 1);

  char *header = (char *)malloc(header_len + 1);
  strncpy(header, token, header_len);
  header[header_len] = '\0';

  char *payload = (char *)malloc(payload_len + 1);
  strncpy(payload, dot1 + 1, payload_len);
  payload[payload_len] = '\0';

  char *signature = strndup(dot2 + 1, sig_len);

  char *signing_input = (char *)malloc(header_len + payload_len + 2);
  sprintf(signing_input, "%s.%s", header, payload);

  int expected_sig_len;
  unsigned char *expected_sig_bin =
      jwt_base64url_decode(signature, sig_len, &expected_sig_len);

  unsigned int calculated_sig_len;
  unsigned char *calculated_sig =
      jwt_hmac_sha256(secret, strlen(secret), signing_input,
                      strlen(signing_input), &calculated_sig_len);

  bool sig_valid =
      (calculated_sig_len == expected_sig_len) &&
      (memcmp(calculated_sig, expected_sig_bin, calculated_sig_len) == 0);

  bool exp_valid = true;
  if (sig_valid) {
    int payload_bin_len;
    unsigned char *payload_bin =
        jwt_base64url_decode(payload, payload_len, &payload_bin_len);
    char *payload_str = (char *)malloc(payload_bin_len + 1);
    memcpy(payload_str, payload_bin, payload_bin_len);
    payload_str[payload_bin_len] = '\0';

    cJSON *json_payload = cJSON_Parse(payload_str);
    if (json_payload) {
      cJSON *exp = cJSON_GetObjectItemCaseSensitive(json_payload, "exp");
      if (cJSON_IsNumber(exp)) {
        time_t now = time(nullptr);
        if (now >= (time_t)exp->valuedouble) {
          exp_valid = false;
        }
      }
      cJSON_Delete(json_payload);
    }
    free(payload_bin);
    free(payload_str);
  }

  free(header);
  free(payload);
  free(signature);
  free(signing_input);
  free(expected_sig_bin);
  free(calculated_sig);

  return sig_valid && exp_valid;
}

// Получение payload из токена
cJSON *jwt_get_payload(const char *token) {
  const char *dot1 = strchr(token, '.');
  const char *dot2 = dot1 ? strchr(dot1 + 1, '.') : nullptr;
  if (!dot1 || !dot2) return nullptr;

  size_t payload_len = dot2 - dot1 - 1;
  char *payload = (char *)malloc(payload_len + 1);
  strncpy(payload, dot1 + 1, payload_len);
  payload[payload_len] = '\0';

  int payload_bin_len;
  unsigned char *payload_bin =
      jwt_base64url_decode(payload, payload_len, &payload_bin_len);
  char *payload_str = (char *)malloc(payload_bin_len + 1);
  memcpy(payload_str, payload_bin, payload_bin_len);
  payload_str[payload_bin_len] = '\0';

  cJSON *json_payload = cJSON_Parse(payload_str);

  free(payload);
  free(payload_bin);
  free(payload_str);

  return json_payload;
}

// Получение header из токена
cJSON *jwt_get_header(const char *token) {
  const char *dot1 = strchr(token, '.');
  if (!dot1) return nullptr;

  size_t header_len = dot1 - token;
  char *header = (char *)malloc(header_len + 1);
  strncpy(header, token, header_len);
  header[header_len] = '\0';

  int header_bin_len;
  unsigned char *header_bin =
      jwt_base64url_decode(header, header_len, &header_bin_len);
  char *header_str = (char *)malloc(header_bin_len + 1);
  memcpy(header_str, header_bin, header_bin_len);
  header_str[header_bin_len] = '\0';

  cJSON *json_header = cJSON_Parse(header_str);

  free(header);
  free(header_bin);
  free(header_str);

  return json_header;
}