#include "dtob_internal.h"

/* ------------------------------------------------------------------ */
/*  Byte writer                                                       */
/* ------------------------------------------------------------------ */

void bw_init(ByteWriter *w, int strict_validation)
{
    w->cap = 256;
    w->buf = malloc(w->cap);
    if (!w->buf) { w->cap = 0; }
    w->pos = 0;
    w->error = 0;
    w->strict_validation = strict_validation;
}

void bw_write_byte(ByteWriter *w, uint8_t b)
{
    if (w->pos >= w->cap) {
        w->cap *= 2;
        w->buf = realloc(w->buf, w->cap);
    }
    w->buf[w->pos++] = b;
}

void bw_write_ctrl(ByteWriter *w, uint16_t code)
{
    bw_write_byte(w, 0xC0 | ((code >> 8) & 0x3F));
    bw_write_byte(w, code & 0xFF);
}

/* Write data payload: trit-encode bytes, padded to byte boundary */
void bw_write_data(ByteWriter *w, const uint8_t *bytes, size_t len)
{
    if (len == 0) return;
    uint8_t *encoded = NULL;
    size_t enc_len = trit_encode_padded(bytes, len, &encoded);
    for (size_t i = 0; i < enc_len; i++)
        bw_write_byte(w, encoded[i]);
    free(encoded);
}

/* ------------------------------------------------------------------ */
/*  Encode types header                                               */
/* ------------------------------------------------------------------ */

static void encode_types_header(ByteWriter *w, const DtobTypesHeader *th)
{
    if (!th || th->count == 0) return;

    /* types header: OPEN { OPEN code name opcodes CLOSE_KV|CLOSE_ARR } TYPES_CLOSE */
    bw_write_ctrl(w, DTOB_CODE_OPEN);
    for (size_t i = 0; i < th->count; i++) {
        bw_write_ctrl(w, DTOB_CODE_OPEN);
        bw_write_ctrl(w, th->entries[i].code);
        bw_write_data(w, (const uint8_t *)th->entries[i].name,
                      strlen(th->entries[i].name));
        for (size_t j = 0; j < th->entries[i].n_opcodes; j++)
            bw_write_ctrl(w, th->entries[i].opcodes[j]);
        bw_write_ctrl(w, th->entries[i].is_struct ? DTOB_CODE_ARR_CLOSE
                                                  : DTOB_CODE_KV_CLOSE);
    }
    bw_write_ctrl(w, DTOB_CODE_TYPES_CLOSE);
}

/* ------------------------------------------------------------------ */
/*  Encode value                                                      */
/* ------------------------------------------------------------------ */

/* Determine the narrowest standard width (1,2,4,8) that fits data_len bytes */
static int round_up_width(size_t len)
{
    if (len <= 1) return 1;
    if (len <= 2) return 2;
    if (len <= 4) return 4;
    return 8;
}

/* Write big-endian int data padded/sign-extended to exactly `width` bytes */
static void write_int_padded(ByteWriter *w, const uint8_t *data, size_t data_len,
                              int width, int is_negative)
{
    uint8_t buf[8] = {0};
    uint8_t fill = is_negative ? 0xFF : 0x00;
    int pad = width - (int)data_len;
    if (pad > 0) {
        memset(buf, fill, pad);
        memcpy(buf + pad, data, data_len);
    } else {
        memcpy(buf, data + data_len - width, width);
    }
    bw_write_data(w, buf, width);
}

static void encode_value_t(ByteWriter *w, const DtobValue *v,
                            const DtobTypesHeader *types)
{
    switch ((int)v->type) {
    case JSON_STRING:
        bw_write_ctrl(w, JSON_CODE_STRING);
        bw_write_data(w, v->data, v->data_len);
        break;
    case DTOB_INT: {
        int is_negative = (v->data_len > 0 && (v->data[0] & 0x80));
        int is_unsigned = !is_negative;
        int width = round_up_width(v->data_len);
        int width_index = 0;
        switch (width) { case 2: width_index = 1; break;
                         case 4: width_index = 2; break;
                         case 8: width_index = 3; break; }
        uint8_t code = 8 + (is_unsigned ? 4 : 0) + width_index;
        bw_write_ctrl(w, code);
        write_int_padded(w, v->data, v->data_len, width, is_negative);
        break;
    }
    case DTOB_FLOAT:
        bw_write_ctrl(w, DTOB_CODE_DOUBLE);
        bw_write_data(w, v->data, v->data_len);
        break;
    case DTOB_RAW:
        bw_write_ctrl(w, DTOB_CODE_RAW);
        bw_write_data(w, v->data, v->data_len);
        break;
    case JSON_TRUE:  bw_write_ctrl(w, JSON_CODE_TRUE);  break;
    case JSON_FALSE: bw_write_ctrl(w, JSON_CODE_FALSE); break;
    case JSON_NULL:  bw_write_ctrl(w, JSON_CODE_NULL);  break;
    case DTOB_CUSTOM: {
        DtobCustomType *ct = types ? dtob_types_get(types, v->custom_code) : NULL;
        if (types && !ct) {
            if (w->strict_validation) {
                fprintf(stderr, "dtob FATAL: Custom type opcode %u is not registered in the provided Types Header!\n", v->custom_code);
                w->error = 1;
                break;
            }
        }
        if (ct && w->strict_validation) {
            if (ct->is_struct && v->num_elements != ct->n_opcodes) {
                fprintf(stderr, "dtob FATAL: Struct type %u expects %zu elements but got %zu!\n", v->custom_code, ct->n_opcodes, v->num_elements);
                w->error = 1;
                break;
            } else if (!ct->is_struct && ct->n_opcodes > 1) {
                /* check if inner code mathematically exists in enum values */
                int valid_enum = 0;
                for (size_t i = 0; i < ct->n_opcodes; i++) {
                    if (v->inner_code == ct->opcodes[i]) { valid_enum = 1; break; }
                }
                if (!valid_enum) {
                    fprintf(stderr, "dtob FATAL: Enum inner code %u is invalid for custom type %u!\n", v->inner_code, v->custom_code);
                    w->error = 1;
                    break;
                }
            }
        }
        bw_write_ctrl(w, v->custom_code);
        if (ct && ct->is_struct) {
            /* struct: OPEN opcode1 data1 opcode2 data2 ... CLOSE_ARR */
            bw_write_ctrl(w, DTOB_CODE_OPEN);
            for (size_t i = 0; i < v->num_elements; i++)
                encode_value_t(w, v->elements[i], types);
            bw_write_ctrl(w, DTOB_CODE_ARR_CLOSE);
        } else if (ct && ct->n_opcodes == 0) {
            /* nullable: just the custom code, no inner opcode or data */
        } else if (ct && ct->n_opcodes == 1) {
            /* one-type enum: no inner opcode, just data */
            if (v->data && v->data_len > 0)
                bw_write_data(w, v->data, v->data_len);
        } else if (ct && ct->n_opcodes > 1) {
            /* multi-type enum: write inner opcode + data or struct children */
            if (v->inner_code) {
                bw_write_ctrl(w, v->inner_code);
                DtobCustomType *inner_ct = types ? dtob_types_get(types, v->inner_code) : NULL;
                if (inner_ct && inner_ct->is_struct && v->num_elements > 0) {
                    /* struct variant: OPEN children CLOSE_ARR */
                    bw_write_ctrl(w, DTOB_CODE_OPEN);
                    for (size_t i = 0; i < v->num_elements; i++)
                        encode_value_t(w, v->elements[i], types);
                    bw_write_ctrl(w, DTOB_CODE_ARR_CLOSE);
                } else if (v->data && v->data_len > 0) {
                    bw_write_data(w, v->data, v->data_len);
                }
            } else if (v->data && v->data_len > 0) {
                bw_write_data(w, v->data, v->data_len);
            }
        } else {
            /* no type def (legacy): raw data */
            if (v->data && v->data_len > 0)
                bw_write_data(w, v->data, v->data_len);
        }
        break;
    }
    case DTOB_ARRAY:
        bw_write_ctrl(w, DTOB_CODE_OPEN);
        for (size_t i = 0; i < v->num_elements; i++)
            encode_value_t(w, v->elements[i], types);
        bw_write_ctrl(w, DTOB_CODE_ARR_CLOSE);
        break;
    case DTOB_KV_SET: {
        bw_write_ctrl(w, DTOB_CODE_OPEN);
        for (size_t i = 0; i < v->num_pairs; i++) {
            bw_write_ctrl(w, DTOB_CODE_RAW);
            bw_write_data(w, v->pairs[i].key, v->pairs[i].key_len);
            encode_value_t(w, v->pairs[i].value, types);
        }
        for (size_t i = 0; i < v->num_elements; i++)
            encode_value_t(w, v->elements[i], types);
        bw_write_ctrl(w, DTOB_CODE_KV_CLOSE);
        break;
    }
    }
}

/* --- Public DtobWriter API --- */

void dtob_writer_init(DtobWriter *w)  { bw_init((ByteWriter *)w, 0); }
void dtob_writer_byte(DtobWriter *w, uint8_t b) { bw_write_byte((ByteWriter *)w, b); }
void dtob_writer_ctrl(DtobWriter *w, uint16_t code) { bw_write_ctrl((ByteWriter *)w, code); }
void dtob_writer_data(DtobWriter *w, const uint8_t *bytes, size_t len) { bw_write_data((ByteWriter *)w, bytes, len); }
void dtob_writer_value(DtobWriter *w, const DtobValue *v) { encode_value_t((ByteWriter *)w, v, NULL); }
void dtob_writer_value_typed(DtobWriter *w, const DtobValue *v,
                              const DtobTypesHeader *types) { encode_value_t((ByteWriter *)w, v, types); }

/* Public trit wrappers */
size_t dtob_trit_encode(const uint8_t *bytes, size_t byte_len, uint8_t **out_buf)
{ return trit_encode_padded(bytes, byte_len, out_buf); }
size_t dtob_trit_decode(const uint8_t *buf, size_t buf_len, uint8_t **out_bytes)
{ return trit_decode_padded(buf, buf_len, out_bytes); }

/* --- Public encode API --- */

uint8_t *dtob_encode(const DtobValue *root, size_t *out_len)
{
    return dtob_encode_with_types(root, NULL, 0, out_len);
}

/* Write the body of a collection (items + close) without the leading OPEN.
 * Used when a types header already opened the root collection. */
static void encode_collection_body(ByteWriter *w, const DtobValue *v,
                                    const DtobTypesHeader *types)
{
    if (v->type == DTOB_ARRAY) {
        for (size_t i = 0; i < v->num_elements; i++)
            encode_value_t(w, v->elements[i], types);
        bw_write_ctrl(w, DTOB_CODE_ARR_CLOSE);
    } else if (v->type == DTOB_KV_SET) {
        for (size_t i = 0; i < v->num_pairs; i++) {
            bw_write_ctrl(w, DTOB_CODE_RAW);
            bw_write_data(w, v->pairs[i].key, v->pairs[i].key_len);
            encode_value_t(w, v->pairs[i].value, types);
        }
        for (size_t i = 0; i < v->num_elements; i++)
            encode_value_t(w, v->elements[i], types);
        bw_write_ctrl(w, DTOB_CODE_KV_CLOSE);
    }
}
uint8_t *dtob_encode_with_types(const DtobValue *root,
                                 const DtobTypesHeader *types,
                                 int strict_validation,
                                 size_t *out_len)
{
    ByteWriter w;
    bw_init(&w, strict_validation);

    /* magic number */
    for (int i = 0; i < DTOB_MAGIC_LEN; i++)
        bw_write_byte(&w, (uint8_t)DTOB_MAGIC[i]);

    if (types && types->count > 0) {
        bw_write_ctrl(&w, DTOB_CODE_OPEN);    /* root collection open */
        encode_types_header(&w, types);        /* OPEN <types> TYPES_CLOSE */
        encode_collection_body(&w, root, types); /* items + close */
    } else {
        encode_value_t(&w, root, types);
    }

    if (w.error) {
        free(w.buf);
        *out_len = 0;
        return NULL;
    }

    *out_len = w.pos;
    return w.buf;
}

uint8_t *dtob_encode_chunk(const DtobValue *v,
                           const DtobTypesHeader *types,
                           int strict_validation,
                           size_t *out_len)
{
    ByteWriter w;
    bw_init(&w, strict_validation);
    encode_value_t(&w, v, types);
    
    if (w.error) {
        free(w.buf);
        *out_len = 0;
        return NULL;
    }

    *out_len = w.pos;
    return w.buf;
}
