#include "simdjson_vendor.hh"
#include "simdcbor.hh"
#include <cmath>
#include <cstring>

extern "C" {
#include <caml/mlvalues.h>
#include <caml/memory.h>
#include <caml/alloc.h>
#include <caml/bigarray.h>
#include <caml/custom.h>
#include <caml/fail.h>
}

#ifndef Val_none
#define Val_none Val_int(0)
#endif

using namespace simdjson;

/* ========================================================================= */
/* Custom block accessors                                                    */
/* ========================================================================= */

#define Parser_val(v) (*((dom::parser **) Data_custom_val(v)))
#define Element_val(v) (*((dom::element *) Data_custom_val(v)))
#define Array_val(v) (*((dom::array *) Data_custom_val(v)))
#define Object_val(v) (*((dom::object *) Data_custom_val(v)))
#define Array_iterator_val(v) (*((dom::array::iterator *) Data_custom_val(v)))
#define Object_iterator_val(v) (*((dom::object::iterator *) Data_custom_val(v)))

/* ========================================================================= */
/* Finalizers                                                                */
/* ========================================================================= */

static void parser_finalize(value v) {
    dom::parser *p = Parser_val(v);
    if (p) delete p;
}

/* ========================================================================= */
/* Custom operations                                                         */
/* ========================================================================= */

static struct custom_operations parser_ops = {
    "simdjson.parser",
    parser_finalize,
    custom_compare_default,
    custom_hash_default,
    custom_serialize_default,
    custom_deserialize_default,
    custom_compare_ext_default,
    custom_fixed_length_default
};

static struct custom_operations element_ops = {
    "simdjson.element",
    custom_finalize_default,
    custom_compare_default,
    custom_hash_default,
    custom_serialize_default,
    custom_deserialize_default,
    custom_compare_ext_default,
    custom_fixed_length_default
};

static struct custom_operations array_ops = {
    "simdjson.array",
    custom_finalize_default,
    custom_compare_default,
    custom_hash_default,
    custom_serialize_default,
    custom_deserialize_default,
    custom_compare_ext_default,
    custom_fixed_length_default
};

static struct custom_operations object_ops = {
    "simdjson.object",
    custom_finalize_default,
    custom_compare_default,
    custom_hash_default,
    custom_serialize_default,
    custom_deserialize_default,
    custom_compare_ext_default,
    custom_fixed_length_default
};

static struct custom_operations array_iterator_ops = {
    "simdjson.array_iterator",
    custom_finalize_default,
    custom_compare_default,
    custom_hash_default,
    custom_serialize_default,
    custom_deserialize_default,
    custom_compare_ext_default,
    custom_fixed_length_default
};

static struct custom_operations object_iterator_ops = {
    "simdjson.object_iterator",
    custom_finalize_default,
    custom_compare_default,
    custom_hash_default,
    custom_serialize_default,
    custom_deserialize_default,
    custom_compare_ext_default,
    custom_fixed_length_default
};

/* ========================================================================= */
/* Error code conversion                                                     */
/* ========================================================================= */

/*
 * Maps simdjson error codes to OCaml variant constructors.
 * Must match the conversions in simdjsont_raw.ml.
 */
static int error_code_to_int(error_code ec) {
    switch (ec) {
        case SUCCESS: return 0;
        case CAPACITY: return 1;
        case MEMALLOC: return 2;
        case TAPE_ERROR: return 3;
        case DEPTH_ERROR: return 4;
        case STRING_ERROR: return 5;
        case T_ATOM_ERROR: return 6;
        case F_ATOM_ERROR: return 7;
        case N_ATOM_ERROR: return 8;
        case NUMBER_ERROR: return 9;
        case BIGINT_ERROR: return 10;
        case UTF8_ERROR: return 11;
        case UNINITIALIZED: return 12;
        case EMPTY: return 13;
        case UNESCAPED_CHARS: return 14;
        case UNCLOSED_STRING: return 15;
        case UNSUPPORTED_ARCHITECTURE: return 16;
        case INCORRECT_TYPE: return 17;
        case NUMBER_OUT_OF_RANGE: return 18;
        case INDEX_OUT_OF_BOUNDS: return 19;
        case NO_SUCH_FIELD: return 20;
        case IO_ERROR: return 21;
        case INVALID_JSON_POINTER: return 22;
        case INVALID_URI_FRAGMENT: return 23;
        case UNEXPECTED_ERROR: return 24;
        case PARSER_IN_USE: return 25;
        case OUT_OF_ORDER_ITERATION: return 26;
        case INSUFFICIENT_PADDING: return 27;
        case INCOMPLETE_ARRAY_OR_OBJECT: return 28;
        case SCALAR_DOCUMENT_AS_VALUE: return 29;
        case OUT_OF_BOUNDS: return 30;
        case TRAILING_CONTENT: return 31;
        case OUT_OF_CAPACITY: return 32;
        case UNKNOWN_FIELD: return 33;
        default: return 24; // Map unknown to UNEXPECTED_ERROR
    }
}

/* ========================================================================= */
/* Parser operations                                                         */
/* ========================================================================= */

extern "C" CAMLprim value simdjson_create_parser(value unit) {
    CAMLparam1(unit);
    CAMLlocal1(v);
    
    dom::parser *p = new dom::parser();
    v = caml_alloc_custom(&parser_ops, sizeof(dom::parser *), 
                          sizeof(dom::parser), 10 * sizeof(dom::parser));
    Parser_val(v) = p;
    
    CAMLreturn(v);
}

extern "C" CAMLprim value simdjson_free_parser(value parser) {
    CAMLparam1(parser);
    
    dom::parser *p = Parser_val(parser);
    if (p) {
        delete p;
        Parser_val(parser) = nullptr;
    }
    
    CAMLreturn(Val_unit);
}

/*
 * Parse JSON from a bigarray buffer.
 * Returns a tuple: (error_code, element option)
 * If error_code is SUCCESS (0), element option is Some element.
 * Otherwise, element option is None.
 */
extern "C" CAMLprim value simdjson_parse(value parser, value buf, value len) {
    CAMLparam3(parser, buf, len);
    CAMLlocal3(result, element_val, some_val);
    
    dom::parser *p = Parser_val(parser);
    const uint8_t *data = (const uint8_t *)Caml_ba_data_val(buf);
    size_t length = Long_val(len);
    
    dom::element doc;
    auto error = p->parse(data, length, false).get(doc);
    
    // Allocate result tuple (error_code, element option)
    result = caml_alloc_tuple(2);
    Store_field(result, 0, Val_int(error_code_to_int(error)));
    
    if (error == SUCCESS) {
        // Allocate element
        element_val = caml_alloc_custom(&element_ops, sizeof(dom::element), 0, 1);
        Element_val(element_val) = doc;
        
        // Wrap in Some
        some_val = caml_alloc(1, 0);
        Store_field(some_val, 0, element_val);
        Store_field(result, 1, some_val);
    } else {
        // None
        Store_field(result, 1, Val_none);
    }
    
    CAMLreturn(result);
}

/*
 * Parse JSON from a file.
 * Returns a tuple: (error_code, element option)
 */
extern "C" CAMLprim value simdjson_parse_file(value parser, value path) {
    CAMLparam2(parser, path);
    CAMLlocal3(result, element_val, some_val);
    
    dom::parser *p = Parser_val(parser);
    const char *filepath = String_val(path);
    
    dom::element doc;
    auto error = p->load(filepath).get(doc);
    
    result = caml_alloc_tuple(2);
    Store_field(result, 0, Val_int(error_code_to_int(error)));
    
    if (error == SUCCESS) {
        element_val = caml_alloc_custom(&element_ops, sizeof(dom::element), 0, 1);
        Element_val(element_val) = doc;
        
        some_val = caml_alloc(1, 0);
        Store_field(some_val, 0, element_val);
        Store_field(result, 1, some_val);
    } else {
        Store_field(result, 1, Val_none);
    }
    
    CAMLreturn(result);
}

/* ========================================================================= */
/* Element type                                                              */
/* ========================================================================= */

/*
 * Returns the element type as an integer matching the element_type variant.
 */
extern "C" CAMLprim value simdjson_element_type(value elt) {
    CAMLparam1(elt);
    
    dom::element e = Element_val(elt);
    dom::element_type t = e.type();
    
    int result;
    switch (t) {
        case dom::element_type::ARRAY:  result = 0; break;
        case dom::element_type::OBJECT: result = 1; break;
        case dom::element_type::INT64:  result = 2; break;
        case dom::element_type::UINT64: result = 3; break;
        case dom::element_type::DOUBLE: result = 4; break;
        case dom::element_type::STRING: result = 5; break;
        case dom::element_type::BOOL:   result = 6; break;
        case dom::element_type::NULL_VALUE: result = 7; break;
        default: result = 7; // Treat unknown as null
    }
    
    CAMLreturn(Val_int(result));
}

/* ========================================================================= */
/* Element value accessors                                                   */
/* ========================================================================= */

/*
 * All getters return a tuple: (error_code, value option)
 * On success, error_code is SUCCESS and value option is Some value.
 * On error, error_code indicates the problem and value option is None.
 */

extern "C" CAMLprim value simdjson_get_bool(value elt) {
    CAMLparam1(elt);
    CAMLlocal2(result, some_val);
    
    dom::element e = Element_val(elt);
    bool b;
    auto error = e.get(b);
    
    result = caml_alloc_tuple(2);
    Store_field(result, 0, Val_int(error_code_to_int(error)));
    
    if (error == SUCCESS) {
        some_val = caml_alloc(1, 0);
        Store_field(some_val, 0, Val_bool(b));
        Store_field(result, 1, some_val);
    } else {
        Store_field(result, 1, Val_none);
    }
    
    CAMLreturn(result);
}

extern "C" CAMLprim value simdjson_get_int64(value elt) {
    CAMLparam1(elt);
    CAMLlocal2(result, some_val);
    
    dom::element e = Element_val(elt);
    int64_t i;
    auto error = e.get(i);
    
    result = caml_alloc_tuple(2);
    Store_field(result, 0, Val_int(error_code_to_int(error)));
    
    if (error == SUCCESS) {
        some_val = caml_alloc(1, 0);
        Store_field(some_val, 0, caml_copy_int64(i));
        Store_field(result, 1, some_val);
    } else {
        Store_field(result, 1, Val_none);
    }
    
    CAMLreturn(result);
}

extern "C" CAMLprim value simdjson_get_uint64(value elt) {
    CAMLparam1(elt);
    CAMLlocal2(result, some_val);
    
    dom::element e = Element_val(elt);
    uint64_t u;
    auto error = e.get(u);
    
    result = caml_alloc_tuple(2);
    Store_field(result, 0, Val_int(error_code_to_int(error)));
    
    if (error == SUCCESS) {
        some_val = caml_alloc(1, 0);
        // Store uint64 as int64 - OCaml side will use integers library
        Store_field(some_val, 0, caml_copy_int64((int64_t)u));
        Store_field(result, 1, some_val);
    } else {
        Store_field(result, 1, Val_none);
    }
    
    CAMLreturn(result);
}

extern "C" CAMLprim value simdjson_get_double(value elt) {
    CAMLparam1(elt);
    CAMLlocal2(result, some_val);
    
    dom::element e = Element_val(elt);
    double d;
    auto error = e.get(d);
    
    result = caml_alloc_tuple(2);
    Store_field(result, 0, Val_int(error_code_to_int(error)));
    
    if (error == SUCCESS) {
        some_val = caml_alloc(1, 0);
        Store_field(some_val, 0, caml_copy_double(d));
        Store_field(result, 1, some_val);
    } else {
        Store_field(result, 1, Val_none);
    }
    
    CAMLreturn(result);
}

extern "C" CAMLprim value simdjson_get_string(value elt) {
    CAMLparam1(elt);
    CAMLlocal2(result, some_val);
    
    dom::element e = Element_val(elt);
    std::string_view sv;
    auto error = e.get(sv);
    
    result = caml_alloc_tuple(2);
    Store_field(result, 0, Val_int(error_code_to_int(error)));
    
    if (error == SUCCESS) {
        some_val = caml_alloc(1, 0);
        Store_field(some_val, 0, caml_alloc_initialized_string(sv.length(), sv.data()));
        Store_field(result, 1, some_val);
    } else {
        Store_field(result, 1, Val_none);
    }
    
    CAMLreturn(result);
}

extern "C" CAMLprim value simdjson_get_array(value elt) {
    CAMLparam1(elt);
    CAMLlocal3(result, arr_val, some_val);
    
    dom::element e = Element_val(elt);
    dom::array arr;
    auto error = e.get(arr);
    
    result = caml_alloc_tuple(2);
    Store_field(result, 0, Val_int(error_code_to_int(error)));
    
    if (error == SUCCESS) {
        arr_val = caml_alloc_custom(&array_ops, sizeof(dom::array), 0, 1);
        Array_val(arr_val) = arr;
        
        some_val = caml_alloc(1, 0);
        Store_field(some_val, 0, arr_val);
        Store_field(result, 1, some_val);
    } else {
        Store_field(result, 1, Val_none);
    }
    
    CAMLreturn(result);
}

extern "C" CAMLprim value simdjson_get_object(value elt) {
    CAMLparam1(elt);
    CAMLlocal3(result, obj_val, some_val);
    
    dom::element e = Element_val(elt);
    dom::object obj;
    auto error = e.get(obj);
    
    result = caml_alloc_tuple(2);
    Store_field(result, 0, Val_int(error_code_to_int(error)));
    
    if (error == SUCCESS) {
        obj_val = caml_alloc_custom(&object_ops, sizeof(dom::object), 0, 1);
        Object_val(obj_val) = obj;
        
        some_val = caml_alloc(1, 0);
        Store_field(some_val, 0, obj_val);
        Store_field(result, 1, some_val);
    } else {
        Store_field(result, 1, Val_none);
    }
    
    CAMLreturn(result);
}

/* ========================================================================= */
/* Array operations                                                          */
/* ========================================================================= */

extern "C" CAMLprim value simdjson_array_size(value arr) {
    CAMLparam1(arr);
    dom::array a = Array_val(arr);
    CAMLreturn(Val_long(a.size()));
}

extern "C" CAMLprim value simdjson_array_begin(value arr) {
    CAMLparam1(arr);
    CAMLlocal1(iter_val);
    
    dom::array a = Array_val(arr);
    iter_val = caml_alloc_custom(&array_iterator_ops, sizeof(dom::array::iterator), 0, 1);
    Array_iterator_val(iter_val) = a.begin();
    
    CAMLreturn(iter_val);
}

extern "C" CAMLprim value simdjson_array_end(value arr) {
    CAMLparam1(arr);
    CAMLlocal1(iter_val);
    
    dom::array a = Array_val(arr);
    iter_val = caml_alloc_custom(&array_iterator_ops, sizeof(dom::array::iterator), 0, 1);
    Array_iterator_val(iter_val) = a.end();
    
    CAMLreturn(iter_val);
}

extern "C" CAMLprim value simdjson_array_iter_get(value iter) {
    CAMLparam1(iter);
    CAMLlocal1(elt_val);
    
    dom::array::iterator it = Array_iterator_val(iter);
    elt_val = caml_alloc_custom(&element_ops, sizeof(dom::element), 0, 1);
    Element_val(elt_val) = *it;
    
    CAMLreturn(elt_val);
}

extern "C" CAMLprim value simdjson_array_iter_next(value iter) {
    CAMLparam1(iter);
    ++(Array_iterator_val(iter));
    CAMLreturn(Val_unit);
}

extern "C" CAMLprim value simdjson_array_iter_equal(value iter1, value iter2) {
    CAMLparam2(iter1, iter2);
    dom::array::iterator it1 = Array_iterator_val(iter1);
    dom::array::iterator it2 = Array_iterator_val(iter2);
    CAMLreturn(Val_bool(it1 == it2));
}

/* ========================================================================= */
/* Object operations                                                         */
/* ========================================================================= */

extern "C" CAMLprim value simdjson_object_size(value obj) {
    CAMLparam1(obj);
    dom::object o = Object_val(obj);
    CAMLreturn(Val_long(o.size()));
}

extern "C" CAMLprim value simdjson_object_begin(value obj) {
    CAMLparam1(obj);
    CAMLlocal1(iter_val);
    
    dom::object o = Object_val(obj);
    iter_val = caml_alloc_custom(&object_iterator_ops, sizeof(dom::object::iterator), 0, 1);
    Object_iterator_val(iter_val) = o.begin();
    
    CAMLreturn(iter_val);
}

extern "C" CAMLprim value simdjson_object_end(value obj) {
    CAMLparam1(obj);
    CAMLlocal1(iter_val);
    
    dom::object o = Object_val(obj);
    iter_val = caml_alloc_custom(&object_iterator_ops, sizeof(dom::object::iterator), 0, 1);
    Object_iterator_val(iter_val) = o.end();
    
    CAMLreturn(iter_val);
}

extern "C" CAMLprim value simdjson_object_iter_get(value iter) {
    CAMLparam1(iter);
    CAMLlocal3(result, key_val, elt_val);
    
    dom::object::iterator it = Object_iterator_val(iter);
    dom::key_value_pair kv = *it;
    
    // Allocate tuple (key, element)
    result = caml_alloc_tuple(2);
    
    key_val = caml_alloc_initialized_string(kv.key.length(), kv.key.data());
    Store_field(result, 0, key_val);
    
    elt_val = caml_alloc_custom(&element_ops, sizeof(dom::element), 0, 1);
    Element_val(elt_val) = kv.value;
    Store_field(result, 1, elt_val);
    
    CAMLreturn(result);
}

extern "C" CAMLprim value simdjson_object_iter_next(value iter) {
    CAMLparam1(iter);
    ++(Object_iterator_val(iter));
    CAMLreturn(Val_unit);
}

extern "C" CAMLprim value simdjson_object_iter_equal(value iter1, value iter2) {
    CAMLparam2(iter1, iter2);
    dom::object::iterator it1 = Object_iterator_val(iter1);
    dom::object::iterator it2 = Object_iterator_val(iter2);
    CAMLreturn(Val_bool(it1 == it2));
}

/*
 * Find a field by key. Returns (error_code, element option)
 */
extern "C" CAMLprim value simdjson_object_find(value obj, value key) {
    CAMLparam2(obj, key);
    CAMLlocal3(result, elt_val, some_val);
    
    dom::object o = Object_val(obj);
    const char *k = String_val(key);
    
    dom::element e;
    auto error = o[k].get(e);
    
    result = caml_alloc_tuple(2);
    Store_field(result, 0, Val_int(error_code_to_int(error)));
    
    if (error == SUCCESS) {
        elt_val = caml_alloc_custom(&element_ops, sizeof(dom::element), 0, 1);
        Element_val(elt_val) = e;
        
        some_val = caml_alloc(1, 0);
        Store_field(some_val, 0, elt_val);
        Store_field(result, 1, some_val);
    } else {
        Store_field(result, 1, Val_none);
    }
    
    CAMLreturn(result);
}

/* ========================================================================= */
/* JSON Pointer navigation (RFC 6901)                                        */
/* ========================================================================= */

extern "C" CAMLprim value simdjson_at_pointer(value elt, value pointer) {
    CAMLparam2(elt, pointer);
    CAMLlocal3(result, elt_val, some_val);
    
    dom::element e = Element_val(elt);
    const char *ptr = String_val(pointer);
    
    dom::element target;
    auto error = e.at_pointer(ptr).get(target);
    
    result = caml_alloc_tuple(2);
    Store_field(result, 0, Val_int(error_code_to_int(error)));
    
    if (error == SUCCESS) {
        elt_val = caml_alloc_custom(&element_ops, sizeof(dom::element), 0, 1);
        Element_val(elt_val) = target;
        
        some_val = caml_alloc(1, 0);
        Store_field(some_val, 0, elt_val);
        Store_field(result, 1, some_val);
    } else {
        Store_field(result, 1, Val_none);
    }
    
    CAMLreturn(result);
}

/* ========================================================================= */
/* Serialization (element to string)                                         */
/* ========================================================================= */

extern "C" CAMLprim value simdjson_element_to_string(value elt) {
    CAMLparam1(elt);
    CAMLlocal1(result);
    
    dom::element e = Element_val(elt);
    std::string s = minify(e);
    result = caml_alloc_initialized_string(s.length(), s.data());
    
    CAMLreturn(result);
}

/* ========================================================================= */
/* Error message                                                             */
/* ========================================================================= */

extern "C" CAMLprim value simdjson_error_message(value error_code_val) {
    CAMLparam1(error_code_val);
    CAMLlocal1(result);
    
    int code = Int_val(error_code_val);
    error_code ec;
    
    switch (code) {
        case 0: ec = SUCCESS; break;
        case 1: ec = CAPACITY; break;
        case 2: ec = MEMALLOC; break;
        case 3: ec = TAPE_ERROR; break;
        case 4: ec = DEPTH_ERROR; break;
        case 5: ec = STRING_ERROR; break;
        case 6: ec = T_ATOM_ERROR; break;
        case 7: ec = F_ATOM_ERROR; break;
        case 8: ec = N_ATOM_ERROR; break;
        case 9: ec = NUMBER_ERROR; break;
        case 10: ec = BIGINT_ERROR; break;
        case 11: ec = UTF8_ERROR; break;
        case 12: ec = UNINITIALIZED; break;
        case 13: ec = EMPTY; break;
        case 14: ec = UNESCAPED_CHARS; break;
        case 15: ec = UNCLOSED_STRING; break;
        case 16: ec = UNSUPPORTED_ARCHITECTURE; break;
        case 17: ec = INCORRECT_TYPE; break;
        case 18: ec = NUMBER_OUT_OF_RANGE; break;
        case 19: ec = INDEX_OUT_OF_BOUNDS; break;
        case 20: ec = NO_SUCH_FIELD; break;
        case 21: ec = IO_ERROR; break;
        case 22: ec = INVALID_JSON_POINTER; break;
        case 23: ec = INVALID_URI_FRAGMENT; break;
        case 24: ec = UNEXPECTED_ERROR; break;
        case 25: ec = PARSER_IN_USE; break;
        case 26: ec = OUT_OF_ORDER_ITERATION; break;
        case 27: ec = INSUFFICIENT_PADDING; break;
        case 28: ec = INCOMPLETE_ARRAY_OR_OBJECT; break;
        case 29: ec = SCALAR_DOCUMENT_AS_VALUE; break;
        case 30: ec = OUT_OF_BOUNDS; break;
        case 31: ec = TRAILING_CONTENT; break;
        case 32: ec = OUT_OF_CAPACITY; break;
        case 33: ec = UNKNOWN_FIELD; break;
        default: ec = UNEXPECTED_ERROR; break;
    }
    
    const char *msg = error_message(ec);
    result = caml_copy_string(msg);
    
    CAMLreturn(result);
}

/* ========================================================================= */
/* Padding constant                                                          */
/* ========================================================================= */

extern "C" CAMLprim value simdjson_padding(value unit) {
    CAMLparam1(unit);
    CAMLreturn(Val_int(SIMDJSON_PADDING));
}

/* ========================================================================= */
/* Bigarray <-> string/bytes memcpy blits (used by Simdjsont_obuf)           */
/* ========================================================================= */

/* memcpy [len] bytes from string [src] at [srcoff] into bigarray [ba] at
 * [dstoff]. No allocation is performed, so the GC cannot move the arguments. */
extern "C" CAMLprim value simdjson_blit_string_to_ba(value src, value srcoff,
                                                      value ba, value dstoff,
                                                      value len) {
    char *dst = (char *)Caml_ba_data_val(ba);
    memcpy(dst + Long_val(dstoff), String_val(src) + Long_val(srcoff),
           Long_val(len));
    return Val_unit;
}

/* memcpy [len] bytes from bigarray [ba] at [srcoff] into bytes [dst] at
 * [dstoff]. */
extern "C" CAMLprim value simdjson_blit_ba_to_bytes(value ba, value srcoff,
                                                     value dst, value dstoff,
                                                     value len) {
    const char *src = (const char *)Caml_ba_data_val(ba);
    memcpy(Bytes_val(dst) + Long_val(dstoff), src + Long_val(srcoff),
           Long_val(len));
    return Val_unit;
}

/* ========================================================================= */
/* Fast double-to-string for encoding                                        */
/* ========================================================================= */

/*
 * Fast double to string conversion using snprintf.
 * Returns JSON-compatible string representation.
 * NaN and Infinity are converted to "null" per JSON spec.
 */
extern "C" CAMLprim value simdjson_format_double(value v_double) {
    CAMLparam1(v_double);
    CAMLlocal1(result);
    
    double d = Double_val(v_double);
    char buf[32];
    int len;
    
    // Handle special values per JSON spec
    if (std::isnan(d) || std::isinf(d)) {
        result = caml_alloc_initialized_string(4, "null");
        CAMLreturn(result);
    }
    
    // Use %.17g for round-trip precision
    len = snprintf(buf, sizeof(buf), "%.17g", d);
    
    result = caml_alloc_initialized_string(len, buf);
    CAMLreturn(result);
}

/*
 * Write double directly to a buffer at given offset.
 * Returns the number of bytes written.
 * Buffer must have at least 32 bytes available.
 */
extern "C" CAMLprim value simdjson_format_double_to_bytes(value v_bytes, value v_offset, value v_double) {
    CAMLparam3(v_bytes, v_offset, v_double);
    
    char *buf = (char *)Bytes_val(v_bytes);
    size_t offset = Long_val(v_offset);
    double d = Double_val(v_double);
    int len;
    
    // Handle special values per JSON spec
    if (std::isnan(d) || std::isinf(d)) {
        memcpy(buf + offset, "null", 4);
        CAMLreturn(Val_int(4));
    }
    
    len = snprintf(buf + offset, 32, "%.17g", d);
    CAMLreturn(Val_int(len));
}

/* ========================================================================= */
/* Document Stream (NDJSON/parse_many)                                       */
/* ========================================================================= */

/*
 * DocumentStreamState holds the state for iterating over multiple JSON documents.
 * It owns the document_stream and tracks iteration state.
 * 
 * Lifetime: The parser must outlive the stream. We ensure this by having OCaml
 * hold references to both, with the stream referencing the parser value.
 */
struct DocumentStreamState {
    dom::document_stream stream;
    dom::document_stream::iterator current;
    dom::document_stream::iterator end_iter;
    bool started;
    bool finished;
    size_t doc_index;
    
    DocumentStreamState(dom::document_stream&& s) 
        : stream(std::move(s))
        , current()
        , end_iter()
        , started(false)
        , finished(false)
        , doc_index(0) {}
};

#define DocStream_val(v) (*((DocumentStreamState **) Data_custom_val(v)))

static void doc_stream_finalize(value v) {
    DocumentStreamState *state = DocStream_val(v);
    if (state) {
        delete state;
    }
}

static struct custom_operations doc_stream_ops = {
    "simdjson.document_stream",
    doc_stream_finalize,
    custom_compare_default,
    custom_hash_default,
    custom_serialize_default,
    custom_deserialize_default,
    custom_compare_ext_default,
    custom_fixed_length_default
};

/*
 * Create a document stream for parsing multiple JSON documents.
 * 
 * Arguments:
 *   parser - the parser to use (must remain alive during iteration)
 *   buf - bigarray buffer containing NDJSON data (must have SIMDJSON_PADDING)
 *   len - length of actual data (not including padding)
 *   batch_size - maximum size of a single document
 * 
 * Returns: (error_code, stream option)
 */
extern "C" CAMLprim value simdjson_parse_many(value parser, value buf, value len, value batch_size) {
    CAMLparam4(parser, buf, len, batch_size);
    CAMLlocal3(result, stream_val, some_val);
    
    dom::parser *p = Parser_val(parser);
    const uint8_t *data = (const uint8_t *)Caml_ba_data_val(buf);
    size_t length = Long_val(len);
    size_t bsize = Long_val(batch_size);
    
    dom::document_stream stream;
    auto error = p->parse_many(data, length, bsize).get(stream);
    
    result = caml_alloc_tuple(2);
    Store_field(result, 0, Val_int(error_code_to_int(error)));
    
    if (error == SUCCESS) {
        DocumentStreamState *state = new DocumentStreamState(std::move(stream));
        
        stream_val = caml_alloc_custom(&doc_stream_ops, sizeof(DocumentStreamState *),
                                       sizeof(DocumentStreamState), 10 * sizeof(DocumentStreamState));
        DocStream_val(stream_val) = state;
        
        some_val = caml_alloc(1, 0);
        Store_field(some_val, 0, stream_val);
        Store_field(result, 1, some_val);
    } else {
        Store_field(result, 1, Val_none);
    }
    
    CAMLreturn(result);
}

/*
 * Get the next document from the stream.
 * 
 * Returns a tuple: (status, error_code, element option, byte_offset)
 * 
 * status:
 *   0 = has document (error_code and element are valid)
 *   1 = end of stream
 * 
 * When status=0:
 *   - If error_code is SUCCESS, element option is Some element
 *   - If error_code is not SUCCESS, element option is None (parse error for this doc)
 */
extern "C" CAMLprim value simdjson_stream_next(value stream) {
    CAMLparam1(stream);
    CAMLlocal4(result, element_val, some_val, offset_val);
    
    DocumentStreamState *state = DocStream_val(stream);

    // Initialize iteration on first call; on subsequent calls advance past the
    // document handed out previously. Advancing must happen here (not after
    // extracting the element) because simdjson reuses one internal buffer per
    // document: `++current` parses the next document over the current one, so
    // the element we return stays valid only until the next call.
    if (!state->started) {
        state->current = state->stream.begin();
        state->end_iter = state->stream.end();
        state->started = true;
    } else if (!state->finished) {
        ++state->current;
    }

    // Check if we're at the end (use != since simdjson only provides operator!=)
    if (state->finished || !(state->current != state->end_iter)) {
        state->finished = true;
        // Return (1, 0, None, 0) - end of stream
        result = caml_alloc_tuple(4);
        Store_field(result, 0, Val_int(1));  // status = end
        Store_field(result, 1, Val_int(0));  // error_code = SUCCESS (unused)
        Store_field(result, 2, Val_none);    // element = None
        Store_field(result, 3, Val_long(0)); // byte_offset = 0
        CAMLreturn(result);
    }

    auto doc_result = *state->current;
    size_t byte_offset = state->current.current_index();
    ++state->doc_index;

    // Build result
    result = caml_alloc_tuple(4);
    Store_field(result, 0, Val_int(0));  // status = has document
    Store_field(result, 3, Val_long(byte_offset));
    
    dom::element doc;
    auto error = doc_result.get(doc);
    Store_field(result, 1, Val_int(error_code_to_int(error)));
    
    if (error == SUCCESS) {
        element_val = caml_alloc_custom(&element_ops, sizeof(dom::element), 0, 1);
        Element_val(element_val) = doc;
        
        some_val = caml_alloc(1, 0);
        Store_field(some_val, 0, element_val);
        Store_field(result, 2, some_val);
    } else {
        Store_field(result, 2, Val_none);
    }
    
    CAMLreturn(result);
}

/*
 * Get the number of truncated bytes at the end of the stream.
 * Should be called after iteration is complete.
 */
extern "C" CAMLprim value simdjson_stream_truncated_bytes(value stream) {
    CAMLparam1(stream);
    DocumentStreamState *state = DocStream_val(stream);
    CAMLreturn(Val_long(state->stream.truncated_bytes()));
}

/*
 * Get the current document index (0-based count of documents seen).
 */
extern "C" CAMLprim value simdjson_stream_doc_index(value stream) {
    CAMLparam1(stream);
    DocumentStreamState *state = DocStream_val(stream);
    CAMLreturn(Val_long(state->doc_index));
}

/*
 * Check if the stream has finished iterating.
 */
extern "C" CAMLprim value simdjson_stream_is_finished(value stream) {
    CAMLparam1(stream);
    DocumentStreamState *state = DocStream_val(stream);
    CAMLreturn(Val_bool(state->finished));
}

/*
 * Get the total size in bytes of the input.
 */
extern "C" CAMLprim value simdjson_stream_size_in_bytes(value stream) {
    CAMLparam1(stream);
    DocumentStreamState *state = DocStream_val(stream);
    CAMLreturn(Val_long(state->stream.size_in_bytes()));
}

/* ========================================================================= */
/* CBOR operations                                                           */
/* ========================================================================= */

extern "C" CAMLprim value simdjson_parse_cbor(value parser, value buf, value len) {
    CAMLparam3(parser, buf, len);
    CAMLlocal3(result, element_val, some_val);
    
    dom::parser *p = Parser_val(parser);
    const uint8_t *data = (const uint8_t *)Caml_ba_data_val(buf);
    size_t length = Long_val(len);
    
    size_t bytes_read = 0;
    auto error = simdcbor::parse(data, length, *p, bytes_read);
    
    result = caml_alloc_tuple(2);
    Store_field(result, 0, Val_int(error_code_to_int(error)));
    
    if (error == SUCCESS) {
        element_val = caml_alloc_custom(&element_ops, sizeof(dom::element), 0, 1);
        Element_val(element_val) = simdcbor::get_root(*p);
        
        some_val = caml_alloc(1, 0);
        Store_field(some_val, 0, element_val);
        Store_field(result, 1, some_val);
    } else {
        Store_field(result, 1, Val_none);
    }
    
    CAMLreturn(result);
}

extern "C" CAMLprim value simdjson_cbor_next(value parser, value buf, value len, value offset_val) {
    CAMLparam4(parser, buf, len, offset_val);
    CAMLlocal3(result, element_val, some_val);
    
    dom::parser *p = Parser_val(parser);
    const uint8_t *data = (const uint8_t *)Caml_ba_data_val(buf);
    size_t length = Long_val(len);
    size_t offset = Long_val(offset_val);
    
    if (offset >= length) {
        result = caml_alloc_tuple(3);
        Store_field(result, 0, Val_int(error_code_to_int(EMPTY)));
        Store_field(result, 1, Val_none);
        Store_field(result, 2, Val_long(offset));
        CAMLreturn(result);
    }
    
    size_t bytes_read = 0;
    auto error = simdcbor::parse(data + offset, length - offset, *p, bytes_read);
    
    result = caml_alloc_tuple(3);
    Store_field(result, 0, Val_int(error_code_to_int(error)));
    
    if (error == SUCCESS) {
        element_val = caml_alloc_custom(&element_ops, sizeof(dom::element), 0, 1);
        Element_val(element_val) = simdcbor::get_root(*p);
        
        some_val = caml_alloc(1, 0);
        Store_field(some_val, 0, element_val);
        Store_field(result, 1, some_val);
        Store_field(result, 2, Val_long(offset + bytes_read));
    } else {
        Store_field(result, 1, Val_none);
        Store_field(result, 2, Val_long(offset));
    }
    
    CAMLreturn(result);
}