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

static uint16_t varint_encode(uint8_t* out, uint16_t val) {
    if (val < 0x80) {
        if (out) out[0] = (uint8_t)val;
        return 1;
    }
    if (out) {
        out[0] = (uint8_t)(0x80 | (val >> 8));
        out[1] = (uint8_t)(val & 0xFF);
    }
    return 2;
}

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

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

static void set_len_field(uint8_t* data, const struct SerializeField* f, size_t len) {
    if (!f || !data || !f->has_len) 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;
    }
}

static size_t fixed_field_size(const struct SerializeSchema* schema, size_t idx) {
    if (idx + 1 < schema->field_count) {
        return schema->fields[idx + 1].offset - schema->fields[idx].offset;
    }
    return 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 if (f->elem_size == 1) {
            const char* str = *(const char**)(data + f->offset);
            size_t len = f->has_len ? get_len_field(data, f) : (str ? strlen(str) : 0);
            if (len > f->max_size) len = f->max_size;
            total += varint_encode(NULL, (uint16_t)len);
            total += len + 1;
        } else {
            size_t len = f->has_len ? get_len_field(data, f) : f->max_size;
            if (len > f->max_size) len = f->max_size;
            total += varint_encode(NULL, (uint16_t)len);
            total += 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);
            memcpy(ptr, fld, sz);
            ptr += sz;
        } else if (f->elem_size == 1) {
            const char* str = *(const char**)fld;
            size_t len = f->has_len ? get_len_field(data, f) : (str ? strlen(str) : 0);
            if (len > f->max_size) len = f->max_size;
            ptr += varint_encode(ptr, (uint16_t)len);
            if (len > 0) {
                memcpy(ptr, str, len);
                ptr += len;
            }
            *ptr++ = 0;
        } else {
            const uint8_t* arr = *(const uint8_t**)fld;
            size_t len = f->has_len ? get_len_field(data, f) : f->max_size;
            if (len > f->max_size) len = f->max_size;
            ptr += varint_encode(ptr, (uint16_t)len);
            if (arr && len > 0) {
                memcpy(ptr, arr, len * f->elem_size);
                ptr += 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) return SERIALIZE_ERR_NULL;
    
    void* obj = u_malloc(schema->struct_size);
    if (!obj) {
        *structure = NULL;
        return SERIALIZE_ERR_BUF;
    }
    
    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) {
                u_free(obj);
                return SERIALIZE_ERR_BUF;
            }
            memcpy(data + f->offset, ptr, sz);
            ptr += sz;
        } else if (f->elem_size == 1) {
            uint16_t len = varint_decode(ptr, end - ptr, &consumed);
            if (len > f->max_size) len = f->max_size;
            ptr += consumed;
            size_t copy_len = len < (size_t)(end - ptr) ? len : (size_t)(end - ptr);
            
            char* str = NULL;
            if (len > 0) {
                str = u_malloc(len + 1);
                if (str) {
                    if (copy_len > 0) memcpy(str, ptr, copy_len);
                    str[copy_len] = 0;
                }
            } else {
                str = u_malloc(1);
                if (str) str[0] = 0;
            }
            *(char**)(data + f->offset) = str;
            if (f->has_len) {
                set_len_field(data, f, len);
            }
            ptr += len + 1;
        } else {
            uint16_t len = varint_decode(ptr, end - ptr, &consumed);
            if (len > f->max_size) len = f->max_size;
            ptr += consumed;
            size_t arr_sz = len * f->elem_size;
            
            uint8_t* arr = NULL;
            if (arr_sz > 0) {
                arr = u_malloc(arr_sz);
                if (arr && ptr + arr_sz <= end) {
                    memcpy(arr, ptr, arr_sz);
                }
            }
            *(uint8_t**)(data + f->offset) = arr;
            if (f->has_len) {
                set_len_field(data, f, len);
            }
            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;
}