/*
 * v1984 CLI — config and key storage
 *
 * Config file format (binary):
 *   [4 bytes]  magic "V19\x01"
 *   [16 bytes] HMAC-SHA256 signature (truncated)
 *   [12 bytes] AES-GCM nonce
 *   [N bytes]  AES-128-GCM ciphertext (encrypted JSON config)
 *   [16 bytes] AES-GCM tag (appended by GCM)
 *
 * Encryption key derived from a constant embedded in the binary.
 * This is an inconvenience barrier — prevents casual `cat` and
 * copy-paste exfiltration. Real security is vault-side.
 */

#define _POSIX_C_SOURCE 200809L

#include "keystore.h"
#include "util.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <errno.h>

#include "bearssl.h"

#ifdef _WIN32
#include <direct.h>
#define mkdir_p(p) _mkdir(p)
#else
#include <unistd.h>
#define mkdir_p(p) mkdir(p, 0755)
#endif

/* Config encryption key — derived from a string that also appears in the web UI */
static const char CONFIG_SEED[] = "vault1984-l2-";
static const unsigned char CONFIG_MAGIC[4] = { 'V', '1', '9', 0x01 };

/* Derive 16-byte config encryption key from seed */
static void derive_config_key(unsigned char out[16]) {
    br_hkdf_context hkdf;
    br_hkdf_init(&hkdf, &br_sha256_vtable, NULL, 0);
    br_hkdf_inject(&hkdf, (const unsigned char *)CONFIG_SEED, sizeof(CONFIG_SEED) - 1);
    br_hkdf_flip(&hkdf);
    br_hkdf_produce(&hkdf, (const unsigned char *)"key", 3, out, 16);
}

/* Derive 16-byte HMAC key from seed */
static void derive_hmac_key(unsigned char out[16]) {
    br_hkdf_context hkdf;
    br_hkdf_init(&hkdf, &br_sha256_vtable, NULL, 0);
    br_hkdf_inject(&hkdf, (const unsigned char *)CONFIG_SEED, sizeof(CONFIG_SEED) - 1);
    br_hkdf_flip(&hkdf);
    br_hkdf_produce(&hkdf, (const unsigned char *)"hmac", 4, out, 16);
}

/* Compute HMAC-SHA256 signature over data, truncated to 16 bytes */
static void compute_sig(const unsigned char *hmac_key, const unsigned char *data,
                        size_t data_len, unsigned char out[16]) {
    br_hmac_key_context kc;
    br_hmac_key_init(&kc, &br_sha256_vtable, hmac_key, 16);
    br_hmac_context hc;
    br_hmac_init(&hc, &kc, 16);
    br_hmac_update(&hc, data, data_len);
    br_hmac_out(&hc, out);
}

static int get_config_dir(char *buf, size_t len) {
    const char *home = getenv("HOME");
    if (!home) home = getenv("USERPROFILE");
    if (!home) {
        fprintf(stderr, "error: cannot determine home directory\n");
        return -1;
    }
    snprintf(buf, len, "%s/.config/v1984", home);
    return 0;
}

static int ensure_dir(const char *path) {
    struct stat st;
    if (stat(path, &st) == 0 && S_ISDIR(st.st_mode)) return 0;

    /* Create parent (.config) if needed */
    char parent[512];
    snprintf(parent, sizeof(parent), "%s", path);
    char *slash = strrchr(parent, '/');
    if (slash) {
        *slash = '\0';
        struct stat pst;
        if (stat(parent, &pst) != 0) mkdir_p(parent);
    }

    if (mkdir_p(path) != 0 && errno != EEXIST) {
        fprintf(stderr, "error: cannot create %s: %s\n", path, strerror(errno));
        return -1;
    }
    return 0;
}

int keystore_init(const char *vault_url, const char *agent_name,
                  const unsigned char *l2_key, size_t l2_key_len) {
    if (l2_key_len != 16) {
        fprintf(stderr, "error: L2 key must be exactly 16 bytes\n");
        return 1;
    }

    char dir[512];
    if (get_config_dir(dir, sizeof(dir)) != 0) return 1;
    if (ensure_dir(dir) != 0) return 1;

    /* Build JSON config */
    char json[1024];
    char l2_b64[32];
    base64_encode(l2_key, 16, l2_b64, sizeof(l2_b64));

    int json_len = snprintf(json, sizeof(json),
        "{\"vault_url\":\"%s\",\"agent_name\":\"%s\",\"l2_key\":\"%s\"}",
        vault_url, agent_name, l2_b64);

    /* Derive keys */
    unsigned char enc_key[16], hmac_key[16];
    derive_config_key(enc_key);
    derive_hmac_key(hmac_key);

    /* Generate nonce */
    unsigned char nonce[12];
    br_hmac_drbg_context drbg;
    br_hmac_drbg_init(&drbg, &br_sha256_vtable, "v1984-init", 10);
    br_prng_seeder seeder = br_prng_seeder_system(NULL);
    if (seeder) seeder(&drbg.vtable);
    br_hmac_drbg_generate(&drbg, nonce, 12);

    /* Encrypt: AES-128-GCM */
    size_t ct_len = (size_t)json_len + 16; /* ciphertext + tag */
    unsigned char *ct = malloc(ct_len);
    memcpy(ct, json, (size_t)json_len);

    br_aes_ct64_ctr_keys aes_ctx;
    br_aes_ct64_ctr_init(&aes_ctx, enc_key, 16);
    br_gcm_context gcm;
    br_gcm_init(&gcm, &aes_ctx.vtable, br_ghash_ctmul64);
    br_gcm_reset(&gcm, nonce, 12);
    br_gcm_flip(&gcm);
    br_gcm_run(&gcm, 1, ct, (size_t)json_len);
    br_gcm_get_tag(&gcm, ct + json_len);

    /* Build file: magic + sig + nonce + ciphertext+tag */
    size_t file_len = 4 + 16 + 12 + ct_len;
    unsigned char *file_data = malloc(file_len);
    memcpy(file_data, CONFIG_MAGIC, 4);
    /* sig placeholder — filled after */
    memcpy(file_data + 4 + 16, nonce, 12);
    memcpy(file_data + 4 + 16 + 12, ct, ct_len);

    /* Compute HMAC over nonce + ciphertext (everything after sig) */
    compute_sig(hmac_key, file_data + 4 + 16, 12 + ct_len, file_data + 4);

    /* Write file */
    char path[512];
    snprintf(path, sizeof(path), "%s/config", dir);
    FILE *f = fopen(path, "wb");
    if (!f) {
        fprintf(stderr, "error: cannot write %s: %s\n", path, strerror(errno));
        free(ct); free(file_data);
        return 1;
    }
    fwrite(file_data, 1, file_len, f);
    fclose(f);

    free(ct);
    free(file_data);

    /* Clear sensitive data from stack */
    memset(json, 0, sizeof(json));
    memset(enc_key, 0, 16);
    memset(hmac_key, 0, 16);

    fprintf(stderr, "v1984: initialized\n");
    fprintf(stderr, "  vault: %s\n", vault_url);
    fprintf(stderr, "  agent: %s\n", agent_name);
    fprintf(stderr, "  config: %s/config\n", dir);
    return 0;
}

int keystore_load(struct v84_config *cfg) {
    memset(cfg, 0, sizeof(*cfg));

    char dir[512];
    if (get_config_dir(dir, sizeof(dir)) != 0) return -1;

    char path[512];
    snprintf(path, sizeof(path), "%s/config", dir);
    FILE *f = fopen(path, "rb");
    if (!f) {
        fprintf(stderr, "error: not initialized. Run: v1984 init\n");
        return -1;
    }

    fseek(f, 0, SEEK_END);
    long file_len = ftell(f);
    fseek(f, 0, SEEK_SET);

    if (file_len < 4 + 16 + 12 + 16) { /* magic + sig + nonce + min ciphertext */
        fprintf(stderr, "error: config file corrupt\n");
        fclose(f);
        return -1;
    }

    unsigned char *file_data = malloc((size_t)file_len);
    fread(file_data, 1, (size_t)file_len, f);
    fclose(f);

    /* Check magic */
    if (memcmp(file_data, CONFIG_MAGIC, 4) != 0) {
        fprintf(stderr, "error: config file corrupt (bad magic)\n");
        free(file_data);
        return -1;
    }

    /* Derive keys */
    unsigned char enc_key[16], hmac_key[16];
    derive_config_key(enc_key);
    derive_hmac_key(hmac_key);

    /* Verify HMAC */
    unsigned char expected_sig[16];
    size_t payload_len = (size_t)file_len - 4 - 16; /* nonce + ct+tag */
    compute_sig(hmac_key, file_data + 4 + 16, payload_len, expected_sig);

    if (memcmp(file_data + 4, expected_sig, 16) != 0) {
        fprintf(stderr, "error: config file tampered\n");
        free(file_data);
        memset(enc_key, 0, 16);
        memset(hmac_key, 0, 16);
        return -1;
    }

    /* Decrypt */
    unsigned char *nonce = file_data + 4 + 16;
    unsigned char *ct = file_data + 4 + 16 + 12;
    size_t ct_len = (size_t)file_len - 4 - 16 - 12;
    size_t pt_len = ct_len - 16;

    unsigned char *pt = malloc(ct_len);
    memcpy(pt, ct, ct_len);

    br_aes_ct64_ctr_keys aes_ctx;
    br_aes_ct64_ctr_init(&aes_ctx, enc_key, 16);
    br_gcm_context gcm;
    br_gcm_init(&gcm, &aes_ctx.vtable, br_ghash_ctmul64);
    br_gcm_reset(&gcm, nonce, 12);
    br_gcm_flip(&gcm);
    br_gcm_run(&gcm, 0, pt, pt_len);

    if (!br_gcm_check_tag(&gcm, pt + pt_len)) {
        fprintf(stderr, "error: config decryption failed\n");
        free(file_data); free(pt);
        memset(enc_key, 0, 16);
        return -1;
    }
    pt[pt_len] = '\0';

    free(file_data);
    memset(enc_key, 0, 16);
    memset(hmac_key, 0, 16);

    /* Parse JSON manually (avoid cJSON dependency in keystore) */
    /* Extract vault_url */
    char *p = strstr((char *)pt, "\"vault_url\":\"");
    if (p) {
        p += 13;
        char *end = strchr(p, '"');
        if (end) {
            size_t len = (size_t)(end - p);
            if (len >= sizeof(cfg->vault_url)) len = sizeof(cfg->vault_url) - 1;
            memcpy(cfg->vault_url, p, len);
            cfg->vault_url[len] = '\0';
        }
    }

    /* Extract agent_name */
    p = strstr((char *)pt, "\"agent_name\":\"");
    if (p) {
        p += 14;
        char *end = strchr(p, '"');
        if (end) {
            size_t len = (size_t)(end - p);
            if (len >= sizeof(cfg->agent_name)) len = sizeof(cfg->agent_name) - 1;
            memcpy(cfg->agent_name, p, len);
            cfg->agent_name[len] = '\0';
        }
    }

    /* Extract l2_key (base64 → binary) */
    p = strstr((char *)pt, "\"l2_key\":\"");
    if (p) {
        p += 10;
        char *end = strchr(p, '"');
        if (end) {
            *end = '\0';
            size_t key_len = 0;
            base64_decode(p, cfg->l2_key, 16, &key_len);
        }
    }

    /* Clear plaintext from memory */
    memset(pt, 0, pt_len);
    free(pt);

    if (!cfg->vault_url[0]) {
        fprintf(stderr, "error: config missing vault_url\n");
        return -1;
    }

    return 0;
}
