// serialize.c - Universal binary serialization with dynamic allocation
#include "serialize.h"
#include <string.h>

static uint16_t length_encode(uint8_t* out, uint16_t val) {
    if (out) {
        out[0] = (uint8_t)(val >> 8);
        out[1] = (uint8_t)val;
    }
    return 2;
}

static uint16_t length_decode(const uint8_t* in, size_t in_len, size_t* consumed) {
    if (!in || in_len < 2) {
        if (consumed) *consumed = 0;
        return 0;
    }
    if (consumed) *consumed = 2;
    return ((uint16_t)in[0] << 8) | in[1];
}

static size_t get_len_field(const uint8_t* data, const struct SerializeField* f) {
    if (!f || !data || f->len_type == 0) return 0;
    const uint8_t* len_ptr = data + f->len_offset;
    if (f->len_type == 1) {
        return *len_ptr;
    } else if (f->len_type == 2) {
        return *(const uint16_t*)len_ptr;
    } else if (f->len_type == 4) {
        return *(const uint32_t*)len_ptr;
    }
    return 0;
}

static void set_len_field(uint8_t* data, const struct SerializeField* f, size_t len) {
    if (!f || !data || f->len_type == 0) return;
    uint8_t* len_ptr = data + f->len_offset;
    if (f->len_type == 1) {
        *len_ptr = (uint8_t)len;
    } else if (f->len_type == 2) {
        *(uint16_t*)len_ptr = (uint16_t)len;
    } else if (f->len_type == 4) {
        *(uint32_t*)len_ptr = (uint32_t)len;
    }
}

static size_t fixed_field_size(const struct SerializeSchema* schema, size_t idx) {
    if (idx >= schema->field_count) return 0;
    const struct SerializeField* f = &schema->fields[idx];
    if (f->elem_size != 0) return 0;

    size_t next_offset = (idx + 1 < schema->field_count) ?
                         schema->fields[idx + 1].offset :
                         schema->struct_size;

    return (next_offset > f->offset) ? (next_offset - f->offset) : 0;
}

int serialize_encode(const void* structure,
                     const struct SerializeSchema* schema,
                     const uint8_t* header,
                     uint8_t** out_buf,
                     size_t* out_len) {
    if (!schema || !structure || !out_buf || !out_len) return SERIALIZE_ERR_NULL;

    const uint8_t* data = (const uint8_t*)structure;
    size_t total = schema->header_len;

    for (size_t i = 0; i < schema->field_count; i++) {
        const struct SerializeField* f = &schema->fields[i];
        if (f->elem_size == 0) {
            total += fixed_field_size(schema, i);
        } else {
            size_t raw_len;
            if (f->elem_size == 1) { // string
                const char* str = *(const char**)(data + f->offset);
                raw_len = (f->len_type != 0) ? get_len_field(data, f)
                                             : (str ? strlen(str) : 0);
            } else { // array
                const uint8_t* arr = *(const uint8_t**)(data + f->offset);
                raw_len = (f->len_type != 0) ? get_len_field(data, f)
                                             : f->max_size;
                if (raw_len > 0 && !arr) raw_len = 0;
            }
            uint16_t len = (raw_len > f->max_size) ? f->max_size : (uint16_t)raw_len;
            total += length_encode(NULL, len);
            if (f->elem_size == 1) {
                total += (size_t)len + 1; // data + null terminator
            } else {
                total += (size_t)len * f->elem_size;
            }
        }
    }

    uint8_t* buf = u_malloc(total);
    if (!buf) return SERIALIZE_ERR_BUF;

    uint8_t* ptr = buf;
    if (header && schema->header_len > 0) {
        memcpy(ptr, header, schema->header_len);
        ptr += schema->header_len;
    }

    for (size_t i = 0; i < schema->field_count; i++) {
        const struct SerializeField* f = &schema->fields[i];
        const uint8_t* fld = data + f->offset;

        if (f->elem_size == 0) {
            size_t sz = fixed_field_size(schema, i);
            if (sz > 0) {
                memcpy(ptr, fld, sz);
                ptr += sz;
            }
        } else if (f->elem_size == 1) { // string
            const char* str = *(const char**)fld;
            size_t raw_len = (f->len_type != 0) ? get_len_field(data, f)
                                                : (str ? strlen(str) : 0);
            uint16_t len = (raw_len > f->max_size) ? f->max_size : (uint16_t)raw_len;
            if (len > 0 && !str) len = 0;

            ptr += length_encode(ptr, len);
            if (len > 0) {
                memcpy(ptr, str, len);
                ptr += len;
            }
            *ptr++ = '\0';
        } else { // array
            const uint8_t* arr = *(const uint8_t**)fld;
            size_t raw_len = (f->len_type != 0) ? get_len_field(data, f)
                                                : f->max_size;
            uint16_t len = (raw_len > f->max_size) ? f->max_size : (uint16_t)raw_len;
            if (len > 0 && !arr) len = 0;

            ptr += length_encode(ptr, len);
            if (len > 0 && arr) {
                memcpy(ptr, arr, (size_t)len * f->elem_size);
                ptr += (size_t)len * f->elem_size;
            }
        }
    }

    *out_buf = buf;
    *out_len = ptr - buf;
    return SERIALIZE_ERR_OK;
}

int serialize_decode(const uint8_t* in_buf,
                     size_t in_len,
                     const struct SerializeSchema* schema,
                     void** structure) {
    if (!in_buf || !schema || !structure) {
        if (structure) *structure = NULL;
        return SERIALIZE_ERR_NULL;
    }

    void* obj = u_malloc(schema->struct_size);
    if (!obj) {
        *structure = NULL;
        return SERIALIZE_ERR_BUF;
    }
    memset(obj, 0, schema->struct_size);

    const uint8_t* ptr = in_buf;
    const uint8_t* end = in_buf + in_len;
    uint8_t* data = (uint8_t*)obj;
    size_t consumed;

    for (size_t i = 0; i < schema->field_count; i++) {
        const struct SerializeField* f = &schema->fields[i];

        if (f->elem_size == 0) {
            size_t sz = fixed_field_size(schema, i);
            if (ptr + sz > end) {
                serialize_free(schema, &obj);
                *structure = NULL;
                return SERIALIZE_ERR_BUF;
            }
            if (sz > 0) {
                memcpy(data + f->offset, ptr, sz);
                ptr += sz;
            }
        } else if (f->elem_size == 1) { // string
            uint16_t len = length_decode(ptr, end - ptr, &consumed);
            if (consumed == 0) {
                serialize_free(schema, &obj);
                *structure = NULL;
                return SERIALIZE_ERR_BUF;
            }
            ptr += consumed;

            if (len > f->max_size) len = f->max_size;

            size_t needed = (size_t)len + 1;
            if ((size_t)(end - ptr) < needed) {
                serialize_free(schema, &obj);
                *structure = NULL;
                return SERIALIZE_ERR_BUF;
            }

            char* str = NULL;
            if (len == 0) {
                str = u_malloc(1);
                if (str) str[0] = '\0';
            } else {
                str = u_malloc(len + 1);
                if (str) {
                    memcpy(str, ptr, len);
                    str[len] = '\0';
                }
            }

            *(char**)(data + f->offset) = str;
            if (f->len_type != 0) set_len_field(data, f, len);

            if (str == NULL) {
                serialize_free(schema, &obj);
                *structure = NULL;
                return SERIALIZE_ERR_BUF;
            }
            ptr += needed;
        } else { // array
            uint16_t len = length_decode(ptr, end - ptr, &consumed);
            if (consumed == 0) {
                serialize_free(schema, &obj);
                *structure = NULL;
                return SERIALIZE_ERR_BUF;
            }
            ptr += consumed;

            if (len > f->max_size) len = f->max_size;

            size_t arr_sz = (size_t)len * f->elem_size;
            if ((size_t)(end - ptr) < arr_sz) {
                serialize_free(schema, &obj);
                *structure = NULL;
                return SERIALIZE_ERR_BUF;
            }

            uint8_t* arr = NULL;
            if (arr_sz > 0) {
                arr = u_malloc(arr_sz);
                if (arr) memcpy(arr, ptr, arr_sz);
            }

            *(uint8_t**)(data + f->offset) = arr;
            if (f->len_type != 0) set_len_field(data, f, len);

            if (arr_sz > 0 && arr == NULL) {
                serialize_free(schema, &obj);
                *structure = NULL;
                return SERIALIZE_ERR_BUF;
            }
            ptr += arr_sz;
        }
    }

    *structure = obj;
    return SERIALIZE_ERR_OK;
}

void serialize_free(const struct SerializeSchema* schema, void** structure) {
    if (!schema || !structure || !*structure) return;

    uint8_t* data = (uint8_t*)*structure;

    for (size_t i = 0; i < schema->field_count; i++) {
        const struct SerializeField* f = &schema->fields[i];
        if (f->elem_size == 1) {
            char* str = *(char**)(data + f->offset);
            if (str) {
                u_free(str);
                *(char**)(data + f->offset) = NULL;
            }
        } else if (f->elem_size > 1) {
            uint8_t* arr = *(uint8_t**)(data + f->offset);
            if (arr) {
                u_free(arr);
                *(uint8_t**)(data + f->offset) = NULL;
            }
        }
    }

    u_free(*structure);
    *structure = NULL;
}