#include "dtob_internal.h"

DtobValue *ast_make(DtobType type)
{
    DtobValue *v = calloc(1, sizeof(DtobValue));
    if (v) v->type = type;
    return v;
}

void ast_add_element(DtobValue *arr, DtobValue *el)
{
    arr->elements = realloc(arr->elements,
                            (arr->num_elements + 1) * sizeof(DtobValue *));
    arr->elements[arr->num_elements++] = el;
}

void ast_add_pair(DtobValue *kvs, uint8_t *key, size_t key_len,
                  DtobValue *val)
{
    kvs->pairs = realloc(kvs->pairs,
                         (kvs->num_pairs + 1) * sizeof(DtobKVPair));
    kvs->pairs[kvs->num_pairs].key      = key;
    kvs->pairs[kvs->num_pairs].key_len  = key_len;
    kvs->pairs[kvs->num_pairs].value    = val;
    kvs->num_pairs++;
}

/* --- Public constructors --- */

DtobValue *json_string(const char *str, size_t len)
{
    DtobValue *v = ast_make(JSON_STRING);
    if (!v) return NULL;
    v->data = malloc(len ? len : 1);
    if (!v->data) { free(v); return NULL; }
    memcpy(v->data, str, len);
    v->data_len = len;
    return v;
}

DtobValue *dtob_int(int64_t val)
{
    DtobValue *v = ast_make(DTOB_INT);
    if (!v) return NULL;
    /* store as big-endian bytes, trimmed to minimum width */
    uint8_t buf[8];
    uint64_t uval = (uint64_t)val;
    for (int i = 7; i >= 0; i--) {
        buf[i] = uval & 0xFF;
        uval >>= 8;
    }
    /* find first significant byte (keep at least 1) */
    int start = 0;
    if (val >= 0) {
        while (start < 7 && buf[start] == 0x00) start++;
    } else {
        while (start < 7 && buf[start] == 0xFF
               && (buf[start + 1] & 0x80)) start++;
    }
    size_t len = 8 - start;
    v->data = malloc(len);
    memcpy(v->data, buf + start, len);
    v->data_len = len;
    return v;
}

DtobValue *dtob_uint(uint64_t val)
{
    DtobValue *v = ast_make(DTOB_INT);
    if (!v) return NULL;
    /* store as big-endian bytes, trimmed, with leading 0x00 if top bit set
     * so the encoder sees it as unsigned (top bit clear) */
    uint8_t buf[9]; /* extra byte for sign disambiguation */
    buf[0] = 0x00;
    for (int i = 8; i >= 1; i--) {
        buf[i] = val & 0xFF;
        val >>= 8;
    }
    /* trim leading zeros but keep at least 1 byte,
     * and keep buf[0]=0x00 if next byte has top bit set */
    int start = 0;
    while (start < 8 && buf[start] == 0x00 && !(buf[start + 1] & 0x80))
        start++;
    size_t len = 9 - start;
    v->data = malloc(len);
    memcpy(v->data, buf + start, len);
    v->data_len = len;
    return v;
}

DtobValue *dtob_float(double val)
{
    DtobValue *v = ast_make(DTOB_FLOAT);
    if (!v) return NULL;
    v->data = malloc(8);
    memcpy(v->data, &val, 8);
    v->data_len = 8;
    return v;
}

DtobValue *dtob_raw(const uint8_t *data, size_t len)
{
    DtobValue *v = ast_make(DTOB_RAW);
    if (!v) return NULL;
    v->data = malloc(len);
    memcpy(v->data, data, len);
    v->data_len = len;
    return v;
}

DtobValue *json_true(void)  { return ast_make(JSON_TRUE);  }
DtobValue *json_false(void) { return ast_make(JSON_FALSE); }
DtobValue *json_null(void)  { return ast_make(JSON_NULL);  }

DtobValue *dtob_array(void) { return ast_make(DTOB_ARRAY);  }
DtobValue *dtob_kvset(void) { return ast_make(DTOB_KV_SET); }

void dtob_array_push(DtobValue *arr, DtobValue *val)
{
    assert(arr->type == DTOB_ARRAY);
    ast_add_element(arr, val);
}

void dtob_custom_push(DtobValue *custom, DtobValue *val)
{
    assert(custom->type == DTOB_CUSTOM);
    ast_add_element(custom, val);
}

void dtob_kvset_put(DtobValue *kvs, const char *key, DtobValue *val)
{
    assert(kvs->type == DTOB_KV_SET);
    size_t klen = strlen(key);
    uint8_t *kcopy = malloc(klen);
    memcpy(kcopy, key, klen);
    ast_add_pair(kvs, kcopy, klen, val);
}

/* --- value accessors --- */

int64_t dtob_val_to_i64(const DtobValue *v)
{
    if (!v || v->type != DTOB_INT || !v->data || v->data_len == 0) return 0;
    int64_t val = (v->data[0] & 0x80) ? -1 : 0;
    for (size_t i = 0; i < v->data_len; i++)
        val = (val << 8) | v->data[i];
    return val;
}

uint64_t dtob_val_to_u64(const DtobValue *v)
{
    if (!v || v->type != DTOB_INT || !v->data || v->data_len == 0) return 0;
    uint64_t val = 0;
    for (size_t i = 0; i < v->data_len; i++)
        val = (val << 8) | v->data[i];
    return val;
}

size_t dtob_val_to_str(const DtobValue *v, char *out, size_t outsz)
{
    if (!v || !v->data || outsz == 0) { if (outsz) out[0] = '\0'; return 0; }
    if (v->type != DTOB_RAW && v->type != DTOB_CUSTOM) { out[0] = '\0'; return 0; }
    size_t copy = v->data_len < outsz - 1 ? v->data_len : outsz - 1;
    memcpy(out, v->data, copy);
    out[copy] = '\0';
    return copy;
}

/* --- KV-set accessors --- */

DtobValue *dtob_kvset_get(const DtobValue *kvs, const char *key)
{
    if (!kvs || kvs->type != DTOB_KV_SET) return NULL;
    size_t klen = strlen(key);
    for (size_t i = 0; i < kvs->num_pairs; i++) {
        if (kvs->pairs[i].key_len == klen &&
            memcmp(kvs->pairs[i].key, key, klen) == 0)
            return kvs->pairs[i].value;
    }
    return NULL;
}

int64_t dtob_kvset_int(const DtobValue *kvs, const char *key)
{
    DtobValue *v = dtob_kvset_get(kvs, key);
    if (!v || v->type != DTOB_INT || !v->data || v->data_len == 0) return 0;
    /* big-endian, sign-extended */
    int64_t val = (v->data[0] & 0x80) ? -1 : 0;
    for (size_t i = 0; i < v->data_len; i++)
        val = (val << 8) | v->data[i];
    return val;
}

uint64_t dtob_kvset_uint(const DtobValue *kvs, const char *key)
{
    DtobValue *v = dtob_kvset_get(kvs, key);
    if (!v || v->type != DTOB_INT || !v->data || v->data_len == 0) return 0;
    uint64_t val = 0;
    for (size_t i = 0; i < v->data_len; i++)
        val = (val << 8) | v->data[i];
    return val;
}

double dtob_kvset_float(const DtobValue *kvs, const char *key)
{
    DtobValue *v = dtob_kvset_get(kvs, key);
    if (!v || v->type != DTOB_FLOAT || !v->data) return 0.0;
    if (v->data_len == 4) {
        float f; memcpy(&f, v->data, 4); return (double)f;
    }
    if (v->data_len == 8) {
        double d; memcpy(&d, v->data, 8); return d;
    }
    return 0.0;
}

size_t dtob_kvset_str(const DtobValue *kvs, const char *key,
                      char *out, size_t outsz)
{
    DtobValue *v = dtob_kvset_get(kvs, key);
    if (!v || v->type != DTOB_RAW || !v->data) { if (outsz) out[0] = '\0'; return 0; }
    size_t copy = v->data_len < outsz - 1 ? v->data_len : outsz - 1;
    memcpy(out, v->data, copy);
    out[copy] = '\0';
    return copy;
}

const uint8_t *dtob_kvset_raw(const DtobValue *kvs, const char *key,
                              size_t *out_len)
{
    DtobValue *v = dtob_kvset_get(kvs, key);
    if (!v || v->type != DTOB_RAW || !v->data) { if (out_len) *out_len = 0; return NULL; }
    if (out_len) *out_len = v->data_len;
    return v->data;
}

/* --- Pointer cache (hash table) --- */

void ptrcache_init(PtrCache *cache)
{
    memset(cache->buckets, 0, sizeof(cache->buckets));
}

static size_t ptrcache_hash(size_t offset)
{
    return (offset ^ (offset >> 8)) & (PTRCACHE_BUCKETS - 1);
}

DtobValue *ptrcache_lookup(PtrCache *cache, size_t offset)
{
    size_t idx = ptrcache_hash(offset);
    for (PtrCacheEntry *e = cache->buckets[idx]; e; e = e->next)
        if (e->offset == offset) return e->value;
    return NULL;
}

void ptrcache_add(PtrCache *cache, size_t offset, DtobValue *value)
{
    size_t idx = ptrcache_hash(offset);
    PtrCacheEntry *node = malloc(sizeof(PtrCacheEntry));
    if (!node) return;
    node->offset = offset;
    node->value  = value;
    node->next   = cache->buckets[idx];
    cache->buckets[idx] = node;
}

void ptrcache_free(PtrCache *cache)
{
    for (int i = 0; i < PTRCACHE_BUCKETS; i++) {
        PtrCacheEntry *e = cache->buckets[i];
        while (e) {
            PtrCacheEntry *next = e->next;
            free(e);
            e = next;
        }
    }
}

/* --- Deep copy --- */

DtobValue *dtob_deep_copy(const DtobValue *v)
{
    if (!v) return NULL;
    DtobValue *c = ast_make(v->type);
    if (!c) return NULL;
    c->custom_code = v->custom_code;
    c->inner_code = v->inner_code;
    if (v->data && v->data_len > 0) {
        c->data = malloc(v->data_len);
        if (!c->data) { free(c); return NULL; }
        memcpy(c->data, v->data, v->data_len);
        c->data_len = v->data_len;
    }
    for (size_t i = 0; i < v->num_elements; i++)
        ast_add_element(c, dtob_deep_copy(v->elements[i]));
    for (size_t i = 0; i < v->num_pairs; i++) {
        uint8_t *kcopy = malloc(v->pairs[i].key_len);
        memcpy(kcopy, v->pairs[i].key, v->pairs[i].key_len);
        ast_add_pair(c, kcopy, v->pairs[i].key_len,
                     dtob_deep_copy(v->pairs[i].value));
    }
    return c;
}

/* --- Free --- */

void dtob_free(DtobValue *v)
{
    if (!v) return;
    free(v->data);
    for (size_t i = 0; i < v->num_elements; i++)
        dtob_free(v->elements[i]);
    free(v->elements);
    for (size_t i = 0; i < v->num_pairs; i++) {
        free(v->pairs[i].key);
        dtob_free(v->pairs[i].value);
    }
    free(v->pairs);
    free(v);
}

/* --- Print --- */

static void print_bytes_as_string(const uint8_t *data, size_t len)
{
    printf("\"");
    for (size_t i = 0; i < len; i++) {
        if (data[i] >= 32 && data[i] < 127)
            putchar(data[i]);
        else
            printf("\\x%02x", data[i]);
    }
    printf("\"");
}

static void print_indent(int indent)
{
    for (int i = 0; i < indent; i++) printf("  ");
}

void dtob_print(const DtobValue *v, int indent)
{
    if (!v) return;
    print_indent(indent);

    switch ((int)v->type) {
    case JSON_STRING:
        printf("string(");
        print_bytes_as_string(v->data, v->data_len);
        printf(")\n");
        break;
    case DTOB_INT: {
        int is_negative = (v->data_len > 0 && (v->data[0] & 0x80));
        if (is_negative) {
            int64_t val = -1;
            for (size_t i = 0; i < v->data_len; i++)
                val = (val << 8) | v->data[i];
            printf("int(%ld)\n", (long)val);
        } else {
            uint64_t val = 0;
            for (size_t i = 0; i < v->data_len; i++)
                val = (val << 8) | v->data[i];
            printf("uint(%llu)\n", (unsigned long long)val);
        }
        break;
    }
    case DTOB_FLOAT: {
        double d;
        if (v->data_len == 4) {
            float f;
            memcpy(&f, v->data, 4);
            d = (double)f;
        } else {
            memcpy(&d, v->data, 8);
        }
        printf("float(%g)\n", d);
        break;
    }
    case DTOB_RAW:
        printf("raw(%zu bytes)\n", v->data_len);
        break;
    case JSON_TRUE:   printf("true\n");  break;
    case JSON_FALSE:  printf("false\n"); break;
    case JSON_NULL:   printf("null\n");  break;
    case DTOB_ARRAY:
        printf("[\n");
        for (size_t i = 0; i < v->num_elements; i++)
            dtob_print(v->elements[i], indent + 1);
        print_indent(indent);
        printf("]\n");
        break;
    case DTOB_KV_SET:
        printf("{\n");
        for (size_t i = 0; i < v->num_pairs; i++) {
            print_indent(indent + 1);
            print_bytes_as_string(v->pairs[i].key, v->pairs[i].key_len);
            printf(": ");
            /* print value inline if scalar, else newline */
            if (v->pairs[i].value->type == DTOB_ARRAY ||
                v->pairs[i].value->type == DTOB_KV_SET) {
                printf("\n");
                dtob_print(v->pairs[i].value, indent + 2);
            } else {
                dtob_print(v->pairs[i].value, 0);
            }
        }
        print_indent(indent);
        printf("}\n");
        break;
    case DTOB_CUSTOM:
        printf("custom(%d, %zu bytes)\n", v->custom_code, v->data_len);
        break;
    }
}

/* --- Custom type constructor --- */

DtobValue *dtob_custom(uint16_t code, const uint8_t *data, size_t len)
{
    DtobValue *v = ast_make(DTOB_CUSTOM);
    if (!v) return NULL;
    v->custom_code = code;
    if (data && len > 0) {
        v->data = malloc(len);
        memcpy(v->data, data, len);
        v->data_len = len;
    }
    return v;
}

/* --- Types header --- */

void dtob_types_init(DtobTypesHeader *th)
{
    th->cap = 16;
    th->entries = calloc(th->cap, sizeof(DtobCustomType));
    th->count = 0;
}

void dtob_types_cleanup(DtobTypesHeader *th)
{
    for (size_t i = 0; i < th->count; i++)
        free(th->entries[i].name);
    free(th->entries);
    th->entries = NULL;
    th->count = th->cap = 0;
}

int dtob_types_add(DtobTypesHeader *th, uint16_t code, const char *name,
                    const uint16_t *opcodes, size_t n_opcodes)
{
    if (n_opcodes > 15) return -1;
    if (th->count >= th->cap) {
        th->cap = th->cap ? th->cap * 2 : 16;
        th->entries = realloc(th->entries, th->cap * sizeof(DtobCustomType));
    }
    DtobCustomType *e = &th->entries[th->count];
    e->code = code;
    e->name = strdup(name);
    e->n_opcodes = n_opcodes;
    e->is_struct = 0;
    if (opcodes && n_opcodes > 0)
        memcpy(e->opcodes, opcodes, n_opcodes * sizeof(uint16_t));
    th->count++;
    return 0;
}

const char *dtob_types_lookup(const DtobTypesHeader *th, uint16_t code)
{
    if (!th) return NULL;
    for (size_t i = 0; i < th->count; i++) {
        if (th->entries[i].code == code)
            return th->entries[i].name;
    }
    return NULL;
}

DtobCustomType *dtob_types_get(const DtobTypesHeader *th, uint16_t code)
{
    if (!th) return NULL;
    for (size_t i = 0; i < th->count; i++) {
        if (th->entries[i].code == code)
            return (DtobCustomType *)&th->entries[i];
    }
    return NULL;
}

int dtob_opcode_data_size(uint16_t opcode)
{
    if (opcode >= 8 && opcode <= 15)
        return 1 << (opcode & 3); /* 1, 2, 4, 8 */
    if (opcode == DTOB_CODE_FLOAT) return 4;
    if (opcode == DTOB_CODE_DOUBLE) return 8;
    if (opcode >= DTOB_CUSTOM_MIN) return -2; /* custom: caller must check types */
    return -1; /* raw, string = variable */
}

/* --- Print with types --- */

void dtob_print_with_types(const DtobValue *v,
                           const DtobTypesHeader *th, int indent)
{
    if (!v) return;
    if (v->type == DTOB_CUSTOM && th) {
        print_indent(indent);
        const char *name = dtob_types_lookup(th, v->custom_code);
        if (name)
            printf("%s(%zu bytes)\n", name, v->data_len);
        else
            printf("custom(%d, %zu bytes)\n", v->custom_code, v->data_len);
        return;
    }
    dtob_print(v, indent);
}

/* ------------------------------------------------------------------ */
/*  File-level array append                                            */
/* ------------------------------------------------------------------ */

int dtob_array_append_to_file(const char *path, DtobValue *entry)
{
    DtobValue *root = NULL;

    /* read existing file if it exists */
    FILE *f = fopen(path, "rb");
    if (f) {
        fseek(f, 0, SEEK_END);
        long sz = ftell(f);
        fseek(f, 0, SEEK_SET);
        if (sz > 0) {
            uint8_t *buf = malloc((size_t)sz);
            if (buf) {
                fread(buf, 1, (size_t)sz, f);
                root = dtob_decode(buf, (size_t)sz);
                free(buf);
            }
        }
        fclose(f);
    }

    /* if no existing array, create one */
    if (!root || root->type != DTOB_ARRAY) {
        if (root) dtob_free(root);
        root = dtob_array();
    }

    dtob_array_push(root, entry);

    size_t out_len;
    uint8_t *out = dtob_encode(root, &out_len);
    dtob_free(root);
    if (!out) return 1;

    f = fopen(path, "wb");
    if (!f) { free(out); return 1; }
    fwrite(out, 1, out_len, f);
    fclose(f);
    free(out);
    return 0;
}
