#ifndef _DEFAULT_SOURCE
#define _DEFAULT_SOURCE
#endif
#ifdef _WIN32
#ifndef _CRT_RAND_S
#define _CRT_RAND_S
#endif
#endif
#include "cst_realm.h"
#include "cJSON.h"
#include "cst_util.h"
#include "monocypher-ed25519.h"
#include "monocypher.h"
#include "os_compat.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#ifdef _WIN32
#include <io.h>
#else
#include <fcntl.h>
#include <unistd.h>
#endif

static int fill_random(uint8_t *buf, size_t n) {
#ifdef _WIN32
  for (size_t i = 0; i < n; i++) {
    unsigned v = 0;
    if (rand_s(&v) != 0)
      return -1;
    buf[i] = (uint8_t)v;
  }
  return 0;
#else
  FILE *f = fopen("/dev/urandom", "rb");
  if (!f)
    return -1;
  size_t got = fread(buf, 1, n, f);
  fclose(f);
  return got == n ? 0 : -1;
#endif
}

static void hex_encode(const uint8_t *in, size_t n, char *out) {
  static const char *H = "0123456789abcdef";
  for (size_t i = 0; i < n; i++) {
    out[i * 2] = H[in[i] >> 4];
    out[i * 2 + 1] = H[in[i] & 15];
  }
  out[n * 2] = 0;
}

static int hex_nibble(int c) {
  if (c >= '0' && c <= '9')
    return c - '0';
  if (c >= 'a' && c <= 'f')
    return c - 'a' + 10;
  if (c >= 'A' && c <= 'F')
    return c - 'A' + 10;
  return -1;
}

static int hex_decode(const char *s, uint8_t *out, size_t n) {
  if (!s || strlen(s) != n * 2)
    return -1;
  for (size_t i = 0; i < n; i++) {
    int a = hex_nibble(s[i * 2]), b = hex_nibble(s[i * 2 + 1]);
    if (a < 0 || b < 0)
      return -1;
    out[i] = (uint8_t)((a << 4) | b);
  }
  return 0;
}

static int write_hex_file(const char *path, const uint8_t *buf, size_t n) {
  char *hex = malloc(n * 2 + 2);
  if (!hex)
    return -1;
  hex_encode(buf, n, hex);
  hex[n * 2] = '\n';
  hex[n * 2 + 1] = 0;
  char tmp[OS_PATH_MAX + 8];
  if (snprintf(tmp, sizeof(tmp), "%s.tmp", path) >= (int)sizeof(tmp)) {
    free(hex);
    return -1;
  }
#ifdef _WIN32
  FILE *f = fopen(tmp, "wb");
#else
  int fd = open(tmp, O_WRONLY | O_CREAT | O_TRUNC, 0600);
  FILE *f = fd >= 0 ? fdopen(fd, "w") : NULL;
  if (!f && fd >= 0)
    close(fd);
#endif
  if (!f) {
    free(hex);
    return -1;
  }
  int ok = fputs(hex, f) >= 0;
  if (fclose(f) != 0)
    ok = 0;
  free(hex);
  if (!ok) {
    remove(tmp);
    return -1;
  }
  if (os_rename_replace(tmp, path) != 0) {
    remove(tmp);
    return -1;
  }
  (void)os_chmod_private(path);
  return 0;
}

static int read_hex_file(const char *path, uint8_t *buf, size_t n) {
  FILE *f = fopen(path, "rb");
  if (!f)
    return -1;
  char line[256];
  if (!fgets(line, sizeof(line), f)) {
    fclose(f);
    return -1;
  }
  fclose(f);
  size_t len = strlen(line);
  while (len && (line[len - 1] == '\n' || line[len - 1] == '\r'))
    line[--len] = 0;
  return hex_decode(line, buf, n);
}

static void cert_message(const cst_realm_cert_t *c, cst_buf_t *b) {
  char pk[65];
  hex_encode(c->instance_pk, 32, pk);
  cst_buf_puts(b, "CST-REALM-CERT-v1\n");
  cst_buf_printf(b, "instance_id=%s\n", c->instance_id);
  cst_buf_printf(b, "label=%s\n", c->label);
  cst_buf_printf(b, "instance_pk=%s\n", pk);
  cst_buf_printf(b, "not_after=%lld\n", (long long)c->not_after);
}

int cst_realm_dir(char *out, size_t n) {
  char base[OS_PATH_MAX];
  if (os_config_dir(base, sizeof(base)) != 0)
    return -1;
  return os_path_join(out, n, base, "realm");
}

int cst_realm_derive(cst_realm_keys_t *out, const uint8_t seed[32]) {
  if (!out || !seed)
    return -1;
  memcpy(out->seed, seed, 32);
  uint8_t seed_copy[32];
  memcpy(seed_copy, seed, 32);
  crypto_ed25519_key_pair(out->sk, out->pk, seed_copy);
  crypto_wipe(seed_copy, sizeof(seed_copy));
  return 0;
}

int cst_realm_generate(cst_realm_keys_t *out) {
  uint8_t seed[32];
  if (fill_random(seed, sizeof(seed)) != 0)
    return -1;
  int rc = cst_realm_derive(out, seed);
  crypto_wipe(seed, sizeof(seed));
  return rc;
}

int cst_realm_save(const char *dir, const cst_realm_keys_t *k) {
  if (!dir || !k)
    return -1;
  if (os_mkdir_p(dir, 0700) != 0)
    return -1;
  char path[OS_PATH_MAX];
  if (os_path_join(path, sizeof(path), dir, "private.key") != 0)
    return -1;
  if (write_hex_file(path, k->seed, 32) != 0)
    return -1;
  if (os_path_join(path, sizeof(path), dir, "public.key") != 0)
    return -1;
  return write_hex_file(path, k->pk, 32);
}

int cst_realm_load(const char *dir, cst_realm_keys_t *out) {
  if (!dir || !out)
    return -1;
  char path[OS_PATH_MAX];
  uint8_t seed[32];
  if (os_path_join(path, sizeof(path), dir, "private.key") != 0)
    return -1;
  if (read_hex_file(path, seed, 32) != 0)
    return -1;
  int rc = cst_realm_derive(out, seed);
  crypto_wipe(seed, sizeof(seed));
  if (rc != 0)
    return -1;
  uint8_t pk_file[32];
  if (os_path_join(path, sizeof(path), dir, "public.key") == 0 &&
      read_hex_file(path, pk_file, 32) == 0) {
    if (memcmp(pk_file, out->pk, 32) != 0)
      return -1;
  }
  return 0;
}

int cst_realm_fingerprint(const uint8_t pk[32], char out[33]) {
  if (!pk || !out)
    return -1;
  uint8_t h[16];
  crypto_blake2b(h, 16, pk, 32);
  hex_encode(h, 16, out);
  return 0;
}

int cst_realm_derive_password(cst_realm_keys_t *out, const char *password,
                              const uint8_t salt[16]) {
  if (!out || !password || !password[0] || !salt)
    return -1;
  crypto_argon2_config cfg = {
      .algorithm = CRYPTO_ARGON2_ID,
      .nb_blocks = 8192, /* 8 MiB */
      .nb_passes = 3,
      .nb_lanes = 1,
  };
  crypto_argon2_inputs in = {
      .pass = (const uint8_t *)password,
      .salt = salt,
      .pass_size = (uint32_t)strlen(password),
      .salt_size = 16,
  };
  void *work = malloc((size_t)cfg.nb_blocks * 1024);
  if (!work)
    return -1;
  uint8_t seed[32];
  crypto_argon2(seed, 32, work, cfg, in, crypto_argon2_no_extras);
  free(work);
  int rc = cst_realm_derive(out, seed);
  crypto_wipe(seed, sizeof(seed));
  return rc;
}

int cst_realm_is_password_mode(const char *dir) {
  if (!dir)
    return 0;
  char path[OS_PATH_MAX];
  if (os_path_join(path, sizeof(path), dir, "password.salt") != 0)
    return 0;
  FILE *f = fopen(path, "rb");
  if (!f)
    return 0;
  fclose(f);
  return 1;
}

int cst_realm_init_password(const char *dir, const char *password) {
  if (!dir || !password || !password[0])
    return -1;
  if (os_mkdir_p(dir, 0700) != 0)
    return -1;
  uint8_t salt[16];
  if (fill_random(salt, sizeof(salt)) != 0)
    return -1;
  cst_realm_keys_t k;
  if (cst_realm_derive_password(&k, password, salt) != 0)
    return -1;
  char path[OS_PATH_MAX];
  if (os_path_join(path, sizeof(path), dir, "password.salt") != 0)
    return -1;
  if (write_hex_file(path, salt, 16) != 0)
    return -1;
  if (os_path_join(path, sizeof(path), dir, "public.key") != 0)
    return -1;
  if (write_hex_file(path, k.pk, 32) != 0)
    return -1;
  /* Never persist the seed/private key in password mode. */
  if (os_path_join(path, sizeof(path), dir, "mode") != 0)
    return -1;
  FILE *f = fopen(path, "wb");
  if (!f)
    return -1;
  fputs("password\n", f);
  fclose(f);
  crypto_wipe(k.seed, sizeof(k.seed));
  crypto_wipe(k.sk, sizeof(k.sk));
  return 0;
}

int cst_realm_unlock_password(const char *dir, const char *password,
                              cst_realm_keys_t *out) {
  if (!dir || !password || !out)
    return -1;
  char path[OS_PATH_MAX];
  uint8_t salt[16], pk_file[32];
  if (os_path_join(path, sizeof(path), dir, "password.salt") != 0 ||
      read_hex_file(path, salt, 16) != 0)
    return -1;
  if (cst_realm_derive_password(out, password, salt) != 0)
    return -1;
  if (os_path_join(path, sizeof(path), dir, "public.key") != 0 ||
      read_hex_file(path, pk_file, 32) != 0 ||
      memcmp(pk_file, out->pk, 32) != 0) {
    crypto_wipe(out->seed, sizeof(out->seed));
    crypto_wipe(out->sk, sizeof(out->sk));
    memset(out, 0, sizeof(*out));
    return -1; /* wrong password */
  }
  return 0;
}

int cst_realm_enroll(const cst_realm_keys_t *realm, const char *instance_id,
                     const char *label, int64_t not_after, cst_realm_cert_t *cert,
                     uint8_t instance_seed[32], uint8_t instance_sk[64]) {
  if (!realm || !instance_id || !instance_id[0] || !cert || !instance_seed ||
      !instance_sk)
    return -1;
  if (strlen(instance_id) >= sizeof(cert->instance_id))
    return -1;
  memset(cert, 0, sizeof(*cert));
  snprintf(cert->instance_id, sizeof(cert->instance_id), "%s", instance_id);
  snprintf(cert->label, sizeof(cert->label), "%s", label ? label : "");
  cert->not_after = not_after;
  memcpy(cert->realm_pk, realm->pk, 32);
  if (fill_random(instance_seed, 32) != 0)
    return -1;
  uint8_t seed_copy[32];
  memcpy(seed_copy, instance_seed, 32);
  crypto_ed25519_key_pair(instance_sk, cert->instance_pk, seed_copy);
  crypto_wipe(seed_copy, sizeof(seed_copy));
  cst_buf_t msg = {0};
  cert_message(cert, &msg);
  if (!msg.buf) {
    crypto_wipe(instance_seed, 32);
    crypto_wipe(instance_sk, 64);
    return -1;
  }
  crypto_ed25519_sign(cert->sig, realm->sk, (const uint8_t *)msg.buf, msg.len);
  cst_buf_free(&msg);
  return 0;
}

int cst_realm_verify_cert(const uint8_t realm_pk[32], const cst_realm_cert_t *cert,
                          int64_t now_unix) {
  if (!realm_pk || !cert || !cert->instance_id[0])
    return -1;
  if (memcmp(realm_pk, cert->realm_pk, 32) != 0)
    return -1;
  if (cert->not_after != 0 && now_unix > 0 && now_unix > cert->not_after)
    return -1;
  cst_buf_t msg = {0};
  cert_message(cert, &msg);
  if (!msg.buf)
    return -1;
  int ok = crypto_ed25519_check(cert->sig, realm_pk, (const uint8_t *)msg.buf,
                                msg.len) == 0;
  cst_buf_free(&msg);
  return ok ? 0 : -1;
}

char *cst_realm_cert_to_json(const cst_realm_cert_t *cert) {
  if (!cert)
    return NULL;
  char ipk[65], rpk[65], sig[129];
  hex_encode(cert->instance_pk, 32, ipk);
  hex_encode(cert->realm_pk, 32, rpk);
  hex_encode(cert->sig, 64, sig);
  cJSON *o = cJSON_CreateObject();
  cJSON_AddNumberToObject(o, "v", 1);
  cJSON_AddStringToObject(o, "instance_id", cert->instance_id);
  cJSON_AddStringToObject(o, "label", cert->label);
  cJSON_AddStringToObject(o, "instance_pk", ipk);
  cJSON_AddNumberToObject(o, "not_after", (double)cert->not_after);
  cJSON_AddStringToObject(o, "realm_pk", rpk);
  cJSON_AddStringToObject(o, "sig", sig);
  char *s = cJSON_PrintUnformatted(o);
  cJSON_Delete(o);
  return s;
}

int cst_realm_cert_from_json(const char *json, cst_realm_cert_t *out) {
  if (!json || !out)
    return -1;
  memset(out, 0, sizeof(*out));
  cJSON *o = cJSON_Parse(json);
  if (!o)
    return -1;
  const cJSON *id = cJSON_GetObjectItemCaseSensitive(o, "instance_id");
  const cJSON *label = cJSON_GetObjectItemCaseSensitive(o, "label");
  const cJSON *ipk = cJSON_GetObjectItemCaseSensitive(o, "instance_pk");
  const cJSON *na = cJSON_GetObjectItemCaseSensitive(o, "not_after");
  const cJSON *rpk = cJSON_GetObjectItemCaseSensitive(o, "realm_pk");
  const cJSON *sig = cJSON_GetObjectItemCaseSensitive(o, "sig");
  int ok = cJSON_IsString(id) && cJSON_IsString(ipk) && cJSON_IsString(rpk) &&
           cJSON_IsString(sig) && cJSON_IsNumber(na);
  if (ok) {
    snprintf(out->instance_id, sizeof(out->instance_id), "%s", id->valuestring);
    if (cJSON_IsString(label))
      snprintf(out->label, sizeof(out->label), "%s", label->valuestring);
    out->not_after = (int64_t)na->valuedouble;
    ok = hex_decode(ipk->valuestring, out->instance_pk, 32) == 0 &&
         hex_decode(rpk->valuestring, out->realm_pk, 32) == 0 &&
         hex_decode(sig->valuestring, out->sig, 64) == 0;
  }
  cJSON_Delete(o);
  return ok ? 0 : -1;
}

int cst_realm_challenge_sign(const uint8_t instance_sk[64], const void *chal,
                             size_t n, uint8_t sig[64]) {
  if (!instance_sk || !chal || !sig)
    return -1;
  crypto_ed25519_sign(sig, instance_sk, chal, n);
  return 0;
}

int cst_realm_challenge_check(const cst_realm_cert_t *cert, const void *chal,
                              size_t n, const uint8_t sig[64]) {
  if (!cert || !chal || !sig)
    return -1;
  return crypto_ed25519_check(sig, cert->instance_pk, chal, n) == 0 ? 0 : -1;
}

int cst_realm_revoke_add(const char *dir, const char *instance_id) {
  if (!dir || !instance_id || !instance_id[0])
    return -1;
  if (os_mkdir_p(dir, 0700) != 0)
    return -1;
  if (cst_realm_revoked(dir, instance_id))
    return 0;
  char path[OS_PATH_MAX];
  if (os_path_join(path, sizeof(path), dir, "revoked") != 0)
    return -1;
  FILE *f = fopen(path, "ab");
  if (!f)
    return -1;
  int ok = fprintf(f, "%s\n", instance_id) > 0;
  if (fclose(f) != 0)
    ok = 0;
  return ok ? 0 : -1;
}

int cst_realm_revoked(const char *dir, const char *instance_id) {
  if (!dir || !instance_id)
    return 0;
  char path[OS_PATH_MAX];
  if (os_path_join(path, sizeof(path), dir, "revoked") != 0)
    return 0;
  FILE *f = fopen(path, "rb");
  if (!f)
    return 0;
  char line[128];
  int hit = 0;
  while (fgets(line, sizeof(line), f)) {
    size_t len = strlen(line);
    while (len && (line[len - 1] == '\n' || line[len - 1] == '\r'))
      line[--len] = 0;
    if (!strcmp(line, instance_id)) {
      hit = 1;
      break;
    }
  }
  fclose(f);
  return hit;
}

static int instance_dir(const char *instance_id, char *out, size_t n) {
  char realm[OS_PATH_MAX], base[OS_PATH_MAX];
  if (cst_realm_dir(realm, sizeof(realm)) != 0)
    return -1;
  if (os_path_join(base, sizeof(base), realm, "instances") != 0)
    return -1;
  if (os_mkdir_p(base, 0700) != 0)
    return -1;
  return os_path_join(out, n, base, instance_id);
}

int cst_realm_instance_save(const char *instance_id, const cst_realm_cert_t *cert,
                            const uint8_t instance_seed[32]) {
  if (!instance_id || !cert || !instance_seed)
    return -1;
  char dir[OS_PATH_MAX], path[OS_PATH_MAX];
  if (instance_dir(instance_id, dir, sizeof(dir)) != 0 ||
      os_mkdir_p(dir, 0700) != 0)
    return -1;
  char *json = cst_realm_cert_to_json(cert);
  if (!json)
    return -1;
  if (os_path_join(path, sizeof(path), dir, "cert.json") != 0) {
    free(json);
    return -1;
  }
  FILE *f = fopen(path, "wb");
  int ok = 0;
  if (f) {
    ok = fputs(json, f) >= 0 && fputc('\n', f) >= 0;
    if (fclose(f) != 0)
      ok = 0;
  }
  free(json);
  if (!ok)
    return -1;
  (void)os_chmod_private(path);
  if (os_path_join(path, sizeof(path), dir, "seed") != 0)
    return -1;
  return write_hex_file(path, instance_seed, 32);
}

int cst_realm_instance_load(const char *instance_id, cst_realm_cert_t *cert,
                            uint8_t instance_sk[64]) {
  if (!instance_id || !cert || !instance_sk)
    return -1;
  char dir[OS_PATH_MAX], path[OS_PATH_MAX];
  if (instance_dir(instance_id, dir, sizeof(dir)) != 0)
    return -1;
  if (os_path_join(path, sizeof(path), dir, "cert.json") != 0)
    return -1;
  FILE *f = fopen(path, "rb");
  if (!f)
    return -1;
  cst_buf_t b = {0};
  char chunk[4096];
  size_t n;
  while ((n = fread(chunk, 1, sizeof(chunk), f)) && b.len < 256 * 1024)
    cst_buf_append(&b, chunk, n);
  fclose(f);
  int ok = b.buf && cst_realm_cert_from_json(b.buf, cert) == 0;
  cst_buf_free(&b);
  if (!ok)
    return -1;
  uint8_t seed[32];
  if (os_path_join(path, sizeof(path), dir, "seed") != 0 ||
      read_hex_file(path, seed, 32) != 0)
    return -1;
  uint8_t seed_copy[32], pk[32];
  memcpy(seed_copy, seed, 32);
  crypto_ed25519_key_pair(instance_sk, pk, seed_copy);
  crypto_wipe(seed_copy, sizeof(seed_copy));
  crypto_wipe(seed, sizeof(seed));
  if (memcmp(pk, cert->instance_pk, 32) != 0)
    return -1;
  return 0;
}