#include <erl_nif.h>
#include <stdio.h>
#include <string.h>
#include <stdint.h>

#include "ckzg.h"
#include "common/bytes.h"

// Globals for settings
KZGSettings kzg_settings;
int kzg_settings_loaded = 0;

static ERL_NIF_TERM nif_load_trusted_setup(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
    if (argc != 1) return enif_make_badarg(env);

    ErlNifBinary bin;
    if (!enif_inspect_binary(env, argv[0], &bin) || bin.size >= 1024) {
        return enif_make_badarg(env);
    }

    char filename[1024];
    memcpy(filename, bin.data, bin.size);
    filename[bin.size] = '\0';

    FILE* f = fopen(filename, "r");
    if (!f) {
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "file_not_found"));
    }

    if (kzg_settings_loaded) {
        free_trusted_setup(&kzg_settings);
        kzg_settings_loaded = 0;
    }

    C_KZG_RET ret = load_trusted_setup_file(&kzg_settings, f, 0);
    fclose(f);

    if (ret != C_KZG_OK) {
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_int(env, ret));
    }

    kzg_settings_loaded = 1;
    return enif_make_atom(env, "ok");
}

static ERL_NIF_TERM nif_blob_to_kzg_commitment(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
    if (argc != 1) return enif_make_badarg(env);
    if (!kzg_settings_loaded) return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "kzg_settings_not_loaded"));

    ErlNifBinary blob_bin;
    if (!enif_inspect_binary(env, argv[0], &blob_bin) || blob_bin.size != BYTES_PER_BLOB) {
        return enif_make_badarg(env);
    }

    Blob blob;
    memcpy(blob.bytes, blob_bin.data, BYTES_PER_BLOB);

    KZGCommitment c;
    C_KZG_RET ret = blob_to_kzg_commitment(&c, &blob, &kzg_settings);

    if (ret != C_KZG_OK) {
         return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_int(env, ret));
    }

    ErlNifBinary out;
    enif_alloc_binary(BYTES_PER_COMMITMENT, &out);
    memcpy(out.data, c.bytes, BYTES_PER_COMMITMENT);

    return enif_make_tuple2(env, enif_make_atom(env, "ok"), enif_make_binary(env, &out));
}

static ERL_NIF_TERM nif_compute_kzg_proof(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
    if (argc != 2) return enif_make_badarg(env);
     if (!kzg_settings_loaded) return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "kzg_settings_not_loaded"));

    ErlNifBinary blob_bin;
    ErlNifBinary z_bin;

    if (!enif_inspect_binary(env, argv[0], &blob_bin) || blob_bin.size != BYTES_PER_BLOB) return enif_make_badarg(env);
    if (!enif_inspect_binary(env, argv[1], &z_bin) || z_bin.size != BYTES_PER_FIELD_ELEMENT) return enif_make_badarg(env);

    Blob blob;
    memcpy(blob.bytes, blob_bin.data, BYTES_PER_BLOB);
    
    Bytes32 z_bytes;
    memcpy(z_bytes.bytes, z_bin.data, BYTES_PER_FIELD_ELEMENT);

    KZGProof proof;
    Bytes32 y;
    
    C_KZG_RET ret = compute_kzg_proof(&proof, &y, &blob, &z_bytes, &kzg_settings);

    if (ret != C_KZG_OK) {
         return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_int(env, ret));
    }
    
    ErlNifBinary proof_out, y_out;
    enif_alloc_binary(BYTES_PER_PROOF, &proof_out);
    memcpy(proof_out.data, proof.bytes, BYTES_PER_PROOF);
    
    enif_alloc_binary(BYTES_PER_FIELD_ELEMENT, &y_out);
    memcpy(y_out.data, y.bytes, BYTES_PER_FIELD_ELEMENT);

    return enif_make_tuple3(env, enif_make_atom(env, "ok"), enif_make_binary(env, &proof_out), enif_make_binary(env, &y_out));
}

static ERL_NIF_TERM nif_verify_kzg_proof(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
      if (argc != 4) return enif_make_badarg(env);
     if (!kzg_settings_loaded) return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "kzg_settings_not_loaded"));

    ErlNifBinary c_bin, z_bin, y_bin, proof_bin;
    
    if (!enif_inspect_binary(env, argv[0], &c_bin) || c_bin.size != BYTES_PER_COMMITMENT) return enif_make_badarg(env);
    if (!enif_inspect_binary(env, argv[1], &z_bin) || z_bin.size != BYTES_PER_FIELD_ELEMENT) return enif_make_badarg(env);
    if (!enif_inspect_binary(env, argv[2], &y_bin) || y_bin.size != BYTES_PER_FIELD_ELEMENT) return enif_make_badarg(env);
    if (!enif_inspect_binary(env, argv[3], &proof_bin) || proof_bin.size != BYTES_PER_PROOF) return enif_make_badarg(env);

    Bytes48 c_bytes; memcpy(c_bytes.bytes, c_bin.data, BYTES_PER_COMMITMENT);
    Bytes32 z_bytes; memcpy(z_bytes.bytes, z_bin.data, BYTES_PER_FIELD_ELEMENT);
    Bytes32 y_bytes; memcpy(y_bytes.bytes, y_bin.data, BYTES_PER_FIELD_ELEMENT);
    Bytes48 proof_bytes; memcpy(proof_bytes.bytes, proof_bin.data, BYTES_PER_PROOF);

    bool ok;
    C_KZG_RET ret = verify_kzg_proof(&ok, &c_bytes, &z_bytes, &y_bytes, &proof_bytes, &kzg_settings);

     if (ret != C_KZG_OK) {
         return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_int(env, ret));
    }
    
    return ok ? enif_make_atom(env, "true") : enif_make_atom(env, "false");
}

static ERL_NIF_TERM nif_verify_blob_kzg_proof(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
    if (argc != 3) return enif_make_badarg(env);
     if (!kzg_settings_loaded) return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "kzg_settings_not_loaded"));
     
    ErlNifBinary blob_bin, c_bin, proof_bin;
    
    if (!enif_inspect_binary(env, argv[0], &blob_bin) || blob_bin.size != BYTES_PER_BLOB) return enif_make_badarg(env);
    if (!enif_inspect_binary(env, argv[1], &c_bin) || c_bin.size != BYTES_PER_COMMITMENT) return enif_make_badarg(env);
    if (!enif_inspect_binary(env, argv[2], &proof_bin) || proof_bin.size != BYTES_PER_PROOF) return enif_make_badarg(env);

    Blob blob; memcpy(blob.bytes, blob_bin.data, BYTES_PER_BLOB);
    Bytes48 c_bytes; memcpy(c_bytes.bytes, c_bin.data, BYTES_PER_COMMITMENT);
    Bytes48 proof_bytes; memcpy(proof_bytes.bytes, proof_bin.data, BYTES_PER_PROOF);

    bool ok;
    C_KZG_RET ret = verify_blob_kzg_proof(&ok, &blob, &c_bytes, &proof_bytes, &kzg_settings);

     if (ret != C_KZG_OK) {
         return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_int(env, ret));
    }
    
    return ok ? enif_make_atom(env, "true") : enif_make_atom(env, "false");
}

static ERL_NIF_TERM nif_compute_blob_kzg_proof(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
    if (argc != 2) return enif_make_badarg(env);
    if (!kzg_settings_loaded)
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "kzg_settings_not_loaded"));

    ErlNifBinary blob_bin, c_bin;
    if (!enif_inspect_binary(env, argv[0], &blob_bin) || blob_bin.size != BYTES_PER_BLOB) return enif_make_badarg(env);
    if (!enif_inspect_binary(env, argv[1], &c_bin) || c_bin.size != BYTES_PER_COMMITMENT) return enif_make_badarg(env);

    Blob blob; memcpy(blob.bytes, blob_bin.data, BYTES_PER_BLOB);
    Bytes48 c_bytes; memcpy(c_bytes.bytes, c_bin.data, BYTES_PER_COMMITMENT);

    KZGProof proof;
    C_KZG_RET ret = compute_blob_kzg_proof(&proof, &blob, &c_bytes, &kzg_settings);

    if (ret != C_KZG_OK) {
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_int(env, ret));
    }

    ErlNifBinary proof_out;
    enif_alloc_binary(BYTES_PER_PROOF, &proof_out);
    memcpy(proof_out.data, proof.bytes, BYTES_PER_PROOF);

    return enif_make_tuple2(env, enif_make_atom(env, "ok"), enif_make_binary(env, &proof_out));
}

static ERL_NIF_TERM nif_verify_blob_kzg_proof_batch(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
    if (argc != 3) return enif_make_badarg(env);
    if (!kzg_settings_loaded)
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "kzg_settings_not_loaded"));

    ERL_NIF_TERM blobs_list = argv[0];
    ERL_NIF_TERM commitments_list = argv[1];
    ERL_NIF_TERM proofs_list = argv[2];

    unsigned int n_blobs, n_commitments, n_proofs;
    if (!enif_get_list_length(env, blobs_list, &n_blobs)) return enif_make_badarg(env);
    if (!enif_get_list_length(env, commitments_list, &n_commitments)) return enif_make_badarg(env);
    if (!enif_get_list_length(env, proofs_list, &n_proofs)) return enif_make_badarg(env);
    if (n_blobs != n_commitments || n_blobs != n_proofs) return enif_make_badarg(env);

    uint64_t n = (uint64_t)n_blobs;

    Blob* blobs = (Blob*)enif_alloc(sizeof(Blob) * n);
    Bytes48* commitments = (Bytes48*)enif_alloc(sizeof(Bytes48) * n);
    Bytes48* proofs = (Bytes48*)enif_alloc(sizeof(Bytes48) * n);
    if (!blobs || !commitments || !proofs) {
        if (blobs) enif_free(blobs);
        if (commitments) enif_free(commitments);
        if (proofs) enif_free(proofs);
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "alloc_failed"));
    }

    ERL_NIF_TERM head, tail;
    unsigned int i = 0;

    // blobs
    tail = blobs_list;
    while (enif_get_list_cell(env, tail, &head, &tail)) {
        ErlNifBinary bin;
        if (!enif_inspect_binary(env, head, &bin) || bin.size != BYTES_PER_BLOB) {
            enif_free(blobs); enif_free(commitments); enif_free(proofs);
            return enif_make_badarg(env);
        }
        memcpy(blobs[i].bytes, bin.data, BYTES_PER_BLOB);
        i++;
    }

    // commitments
    i = 0;
    tail = commitments_list;
    while (enif_get_list_cell(env, tail, &head, &tail)) {
        ErlNifBinary bin;
        if (!enif_inspect_binary(env, head, &bin) || bin.size != BYTES_PER_COMMITMENT) {
            enif_free(blobs); enif_free(commitments); enif_free(proofs);
            return enif_make_badarg(env);
        }
        memcpy(commitments[i].bytes, bin.data, BYTES_PER_COMMITMENT);
        i++;
    }

    // proofs
    i = 0;
    tail = proofs_list;
    while (enif_get_list_cell(env, tail, &head, &tail)) {
        ErlNifBinary bin;
        if (!enif_inspect_binary(env, head, &bin) || bin.size != BYTES_PER_PROOF) {
            enif_free(blobs); enif_free(commitments); enif_free(proofs);
            return enif_make_badarg(env);
        }
        memcpy(proofs[i].bytes, bin.data, BYTES_PER_PROOF);
        i++;
    }

    bool ok;
    C_KZG_RET ret = verify_blob_kzg_proof_batch(&ok, blobs, commitments, proofs, n, &kzg_settings);

    enif_free(blobs);
    enif_free(commitments);
    enif_free(proofs);

    if (ret != C_KZG_OK) {
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_int(env, ret));
    }

    return enif_make_tuple2(env, enif_make_atom(env, "ok"), ok ? enif_make_atom(env, "true") : enif_make_atom(env, "false"));
}

static ERL_NIF_TERM nif_compute_challenge(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
    if (argc != 2) return enif_make_badarg(env);

    ErlNifBinary blob_bin, c_bin;
    if (!enif_inspect_binary(env, argv[0], &blob_bin) || blob_bin.size != BYTES_PER_BLOB) return enif_make_badarg(env);
    if (!enif_inspect_binary(env, argv[1], &c_bin) || c_bin.size != BYTES_PER_COMMITMENT) return enif_make_badarg(env);

    Blob blob; memcpy(blob.bytes, blob_bin.data, BYTES_PER_BLOB);
    Bytes48 c_bytes; memcpy(c_bytes.bytes, c_bin.data, BYTES_PER_COMMITMENT);

    g1_t commitment;
    if (bytes_to_kzg_commitment(&commitment, &c_bytes) != C_KZG_OK) {
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_int(env, C_KZG_BADARGS));
    }

    fr_t challenge;
    compute_challenge(&challenge, &blob, &commitment);

    Bytes32 out;
    bytes_from_bls_field(&out, &challenge);

    ErlNifBinary out_bin;
    enif_alloc_binary(BYTES_PER_FIELD_ELEMENT, &out_bin);
    memcpy(out_bin.data, out.bytes, BYTES_PER_FIELD_ELEMENT);
    return enif_make_tuple2(env, enif_make_atom(env, "ok"), enif_make_binary(env, &out_bin));
}
  
static ERL_NIF_TERM nif_compute_cells_and_kzg_proofs(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
    if (argc != 1) return enif_make_badarg(env);
    if (!kzg_settings_loaded)
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "kzg_settings_not_loaded"));

    ErlNifBinary blob_bin;
    if (!enif_inspect_binary(env, argv[0], &blob_bin) || blob_bin.size != BYTES_PER_BLOB) return enif_make_badarg(env);

    Blob blob; memcpy(blob.bytes, blob_bin.data, BYTES_PER_BLOB);

    Cell cells[CELLS_PER_EXT_BLOB];
    KZGProof proofs[CELLS_PER_EXT_BLOB];

    C_KZG_RET ret = compute_cells_and_kzg_proofs(cells, proofs, &blob, &kzg_settings);
    if (ret != C_KZG_OK) {
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_int(env, ret));
    }

    ERL_NIF_TERM* cell_terms = (ERL_NIF_TERM*)enif_alloc(sizeof(ERL_NIF_TERM) * CELLS_PER_EXT_BLOB);
    ERL_NIF_TERM* proof_terms = (ERL_NIF_TERM*)enif_alloc(sizeof(ERL_NIF_TERM) * CELLS_PER_EXT_BLOB);
    if (!cell_terms || !proof_terms) {
        if (cell_terms) enif_free(cell_terms);
        if (proof_terms) enif_free(proof_terms);
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "alloc_failed"));
    }

    for (size_t i = 0; i < CELLS_PER_EXT_BLOB; i++) {
        ErlNifBinary cell_out;
        enif_alloc_binary(BYTES_PER_CELL, &cell_out);
        memcpy(cell_out.data, cells[i].bytes, BYTES_PER_CELL);
        cell_terms[i] = enif_make_binary(env, &cell_out);

        ErlNifBinary proof_out;
        enif_alloc_binary(BYTES_PER_PROOF, &proof_out);
        memcpy(proof_out.data, proofs[i].bytes, BYTES_PER_PROOF);
        proof_terms[i] = enif_make_binary(env, &proof_out);
    }

    ERL_NIF_TERM cells_list = enif_make_list_from_array(env, cell_terms, CELLS_PER_EXT_BLOB);
    ERL_NIF_TERM proofs_list = enif_make_list_from_array(env, proof_terms, CELLS_PER_EXT_BLOB);

    enif_free(cell_terms);
    enif_free(proof_terms);

    return enif_make_tuple3(env, enif_make_atom(env, "ok"), cells_list, proofs_list);
}

static ERL_NIF_TERM nif_recover_cells_and_kzg_proofs(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
    if (argc != 2) return enif_make_badarg(env);
    if (!kzg_settings_loaded)
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "kzg_settings_not_loaded"));

    ERL_NIF_TERM indices_list = argv[0];
    ERL_NIF_TERM cells_list = argv[1];

    unsigned int n_indices, n_cells;
    if (!enif_get_list_length(env, indices_list, &n_indices)) return enif_make_badarg(env);
    if (!enif_get_list_length(env, cells_list, &n_cells)) return enif_make_badarg(env);
    if (n_indices != n_cells) return enif_make_badarg(env);
    if (n_cells > CELLS_PER_EXT_BLOB) return enif_make_badarg(env);

    uint64_t* cell_indices = (uint64_t*)enif_alloc(sizeof(uint64_t) * n_cells);
    Cell* cells = (Cell*)enif_alloc(sizeof(Cell) * n_cells);
    if (!cell_indices || !cells) {
        if (cell_indices) enif_free(cell_indices);
        if (cells) enif_free(cells);
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "alloc_failed"));
    }

    ERL_NIF_TERM head, tail;
    unsigned int i = 0;

    tail = indices_list;
    while (enif_get_list_cell(env, tail, &head, &tail)) {
        unsigned long idx;
        if (!enif_get_ulong(env, head, &idx)) {
            enif_free(cell_indices); enif_free(cells);
            return enif_make_badarg(env);
        }
        cell_indices[i++] = (uint64_t)idx;
    }

    i = 0;
    tail = cells_list;
    while (enif_get_list_cell(env, tail, &head, &tail)) {
        ErlNifBinary bin;
        if (!enif_inspect_binary(env, head, &bin) || bin.size != BYTES_PER_CELL) {
            enif_free(cell_indices); enif_free(cells);
            return enif_make_badarg(env);
        }
        memcpy(cells[i].bytes, bin.data, BYTES_PER_CELL);
        i++;
    }

    Cell recovered_cells[CELLS_PER_EXT_BLOB];
    KZGProof recovered_proofs[CELLS_PER_EXT_BLOB];

    C_KZG_RET ret =
        recover_cells_and_kzg_proofs(recovered_cells, recovered_proofs, cell_indices, cells, (uint64_t)n_cells, &kzg_settings);

    enif_free(cell_indices);
    enif_free(cells);

    if (ret != C_KZG_OK) {
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_int(env, ret));
    }

    ERL_NIF_TERM* cell_terms = (ERL_NIF_TERM*)enif_alloc(sizeof(ERL_NIF_TERM) * CELLS_PER_EXT_BLOB);
    ERL_NIF_TERM* proof_terms = (ERL_NIF_TERM*)enif_alloc(sizeof(ERL_NIF_TERM) * CELLS_PER_EXT_BLOB);
    if (!cell_terms || !proof_terms) {
        if (cell_terms) enif_free(cell_terms);
        if (proof_terms) enif_free(proof_terms);
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "alloc_failed"));
    }

    for (size_t j = 0; j < CELLS_PER_EXT_BLOB; j++) {
        ErlNifBinary cell_out;
        enif_alloc_binary(BYTES_PER_CELL, &cell_out);
        memcpy(cell_out.data, recovered_cells[j].bytes, BYTES_PER_CELL);
        cell_terms[j] = enif_make_binary(env, &cell_out);

        ErlNifBinary proof_out;
        enif_alloc_binary(BYTES_PER_PROOF, &proof_out);
        memcpy(proof_out.data, recovered_proofs[j].bytes, BYTES_PER_PROOF);
        proof_terms[j] = enif_make_binary(env, &proof_out);
    }

    ERL_NIF_TERM out_cells_list = enif_make_list_from_array(env, cell_terms, CELLS_PER_EXT_BLOB);
    ERL_NIF_TERM out_proofs_list = enif_make_list_from_array(env, proof_terms, CELLS_PER_EXT_BLOB);

    enif_free(cell_terms);
    enif_free(proof_terms);

    return enif_make_tuple3(env, enif_make_atom(env, "ok"), out_cells_list, out_proofs_list);
}

static ERL_NIF_TERM nif_verify_cell_kzg_proof_batch(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
    if (argc != 4) return enif_make_badarg(env);
    if (!kzg_settings_loaded)
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "kzg_settings_not_loaded"));

    ERL_NIF_TERM commitments_list = argv[0];
    ERL_NIF_TERM indices_list = argv[1];
    ERL_NIF_TERM cells_list = argv[2];
    ERL_NIF_TERM proofs_list = argv[3];

    unsigned int n_commitments, n_indices, n_cells, n_proofs;
    if (!enif_get_list_length(env, commitments_list, &n_commitments)) return enif_make_badarg(env);
    if (!enif_get_list_length(env, indices_list, &n_indices)) return enif_make_badarg(env);
    if (!enif_get_list_length(env, cells_list, &n_cells)) return enif_make_badarg(env);
    if (!enif_get_list_length(env, proofs_list, &n_proofs)) return enif_make_badarg(env);
    if (n_commitments != n_indices || n_commitments != n_cells || n_commitments != n_proofs) return enif_make_badarg(env);

    uint64_t n = (uint64_t)n_cells;

    Bytes48* commitments = (Bytes48*)enif_alloc(sizeof(Bytes48) * n);
    uint64_t* cell_indices = (uint64_t*)enif_alloc(sizeof(uint64_t) * n);
    Cell* cells = (Cell*)enif_alloc(sizeof(Cell) * n);
    Bytes48* proofs = (Bytes48*)enif_alloc(sizeof(Bytes48) * n);
    if (!commitments || !cell_indices || !cells || !proofs) {
        if (commitments) enif_free(commitments);
        if (cell_indices) enif_free(cell_indices);
        if (cells) enif_free(cells);
        if (proofs) enif_free(proofs);
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "alloc_failed"));
    }

    ERL_NIF_TERM head, tail;
    unsigned int i = 0;

    tail = commitments_list;
    while (enif_get_list_cell(env, tail, &head, &tail)) {
        ErlNifBinary bin;
        if (!enif_inspect_binary(env, head, &bin) || bin.size != BYTES_PER_COMMITMENT) {
            enif_free(commitments); enif_free(cell_indices); enif_free(cells); enif_free(proofs);
            return enif_make_badarg(env);
        }
        memcpy(commitments[i].bytes, bin.data, BYTES_PER_COMMITMENT);
        i++;
    }

    i = 0;
    tail = indices_list;
    while (enif_get_list_cell(env, tail, &head, &tail)) {
        unsigned long idx;
        if (!enif_get_ulong(env, head, &idx)) {
            enif_free(commitments); enif_free(cell_indices); enif_free(cells); enif_free(proofs);
            return enif_make_badarg(env);
        }
        cell_indices[i++] = (uint64_t)idx;
    }

    i = 0;
    tail = cells_list;
    while (enif_get_list_cell(env, tail, &head, &tail)) {
        ErlNifBinary bin;
        if (!enif_inspect_binary(env, head, &bin) || bin.size != BYTES_PER_CELL) {
            enif_free(commitments); enif_free(cell_indices); enif_free(cells); enif_free(proofs);
            return enif_make_badarg(env);
        }
        memcpy(cells[i].bytes, bin.data, BYTES_PER_CELL);
        i++;
    }

    i = 0;
    tail = proofs_list;
    while (enif_get_list_cell(env, tail, &head, &tail)) {
        ErlNifBinary bin;
        if (!enif_inspect_binary(env, head, &bin) || bin.size != BYTES_PER_PROOF) {
            enif_free(commitments); enif_free(cell_indices); enif_free(cells); enif_free(proofs);
            return enif_make_badarg(env);
        }
        memcpy(proofs[i].bytes, bin.data, BYTES_PER_PROOF);
        i++;
    }

    bool ok;
    C_KZG_RET ret = verify_cell_kzg_proof_batch(&ok, commitments, cell_indices, cells, proofs, n, &kzg_settings);

    enif_free(commitments);
    enif_free(cell_indices);
    enif_free(cells);
    enif_free(proofs);

    if (ret != C_KZG_OK) {
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_int(env, ret));
    }

    return enif_make_tuple2(env, enif_make_atom(env, "ok"), ok ? enif_make_atom(env, "true") : enif_make_atom(env, "false"));
}

static ERL_NIF_TERM nif_compute_verify_cell_kzg_proof_batch_challenge(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
    if (argc != 5) return enif_make_badarg(env);

    ERL_NIF_TERM commitments_list = argv[0];
    ERL_NIF_TERM commitment_indices_list = argv[1];
    ERL_NIF_TERM cell_indices_list = argv[2];
    ERL_NIF_TERM cells_list = argv[3];
    ERL_NIF_TERM proofs_list = argv[4];

    unsigned int num_commitments, num_commitment_indices, num_cell_indices, num_cells, num_proofs;
    if (!enif_get_list_length(env, commitments_list, &num_commitments)) return enif_make_badarg(env);
    if (!enif_get_list_length(env, commitment_indices_list, &num_commitment_indices)) return enif_make_badarg(env);
    if (!enif_get_list_length(env, cell_indices_list, &num_cell_indices)) return enif_make_badarg(env);
    if (!enif_get_list_length(env, cells_list, &num_cells)) return enif_make_badarg(env);
    if (!enif_get_list_length(env, proofs_list, &num_proofs)) return enif_make_badarg(env);
    if (num_commitment_indices != num_cells || num_cell_indices != num_cells || num_proofs != num_cells) return enif_make_badarg(env);

    Bytes48* commitments = (Bytes48*)enif_alloc(sizeof(Bytes48) * num_commitments);
    uint64_t* commitment_indices = (uint64_t*)enif_alloc(sizeof(uint64_t) * num_cells);
    uint64_t* cell_indices = (uint64_t*)enif_alloc(sizeof(uint64_t) * num_cells);
    Cell* cells = (Cell*)enif_alloc(sizeof(Cell) * num_cells);
    Bytes48* proofs = (Bytes48*)enif_alloc(sizeof(Bytes48) * num_cells);
    if (!commitments || !commitment_indices || !cell_indices || !cells || !proofs) {
        if (commitments) enif_free(commitments);
        if (commitment_indices) enif_free(commitment_indices);
        if (cell_indices) enif_free(cell_indices);
        if (cells) enif_free(cells);
        if (proofs) enif_free(proofs);
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "alloc_failed"));
    }

    ERL_NIF_TERM head, tail;
    unsigned int i = 0;

    tail = commitments_list;
    while (enif_get_list_cell(env, tail, &head, &tail)) {
        ErlNifBinary bin;
        if (!enif_inspect_binary(env, head, &bin) || bin.size != BYTES_PER_COMMITMENT) {
            enif_free(commitments); enif_free(commitment_indices); enif_free(cell_indices); enif_free(cells); enif_free(proofs);
            return enif_make_badarg(env);
        }
        memcpy(commitments[i].bytes, bin.data, BYTES_PER_COMMITMENT);
        i++;
    }

    i = 0;
    tail = commitment_indices_list;
    while (enif_get_list_cell(env, tail, &head, &tail)) {
        unsigned long idx;
        if (!enif_get_ulong(env, head, &idx)) {
            enif_free(commitments); enif_free(commitment_indices); enif_free(cell_indices); enif_free(cells); enif_free(proofs);
            return enif_make_badarg(env);
        }
        commitment_indices[i++] = (uint64_t)idx;
    }

    i = 0;
    tail = cell_indices_list;
    while (enif_get_list_cell(env, tail, &head, &tail)) {
        unsigned long idx;
        if (!enif_get_ulong(env, head, &idx)) {
            enif_free(commitments); enif_free(commitment_indices); enif_free(cell_indices); enif_free(cells); enif_free(proofs);
            return enif_make_badarg(env);
        }
        cell_indices[i++] = (uint64_t)idx;
    }

    i = 0;
    tail = cells_list;
    while (enif_get_list_cell(env, tail, &head, &tail)) {
        ErlNifBinary bin;
        if (!enif_inspect_binary(env, head, &bin) || bin.size != BYTES_PER_CELL) {
            enif_free(commitments); enif_free(commitment_indices); enif_free(cell_indices); enif_free(cells); enif_free(proofs);
            return enif_make_badarg(env);
        }
        memcpy(cells[i].bytes, bin.data, BYTES_PER_CELL);
        i++;
    }

    i = 0;
    tail = proofs_list;
    while (enif_get_list_cell(env, tail, &head, &tail)) {
        ErlNifBinary bin;
        if (!enif_inspect_binary(env, head, &bin) || bin.size != BYTES_PER_PROOF) {
            enif_free(commitments); enif_free(commitment_indices); enif_free(cell_indices); enif_free(cells); enif_free(proofs);
            return enif_make_badarg(env);
        }
        memcpy(proofs[i].bytes, bin.data, BYTES_PER_PROOF);
        i++;
    }

    fr_t challenge;
    C_KZG_RET ret = compute_verify_cell_kzg_proof_batch_challenge(
        &challenge,
        commitments,
        (uint64_t)num_commitments,
        commitment_indices,
        cell_indices,
        cells,
        proofs,
        (uint64_t)num_cells
    );

    enif_free(commitments);
    enif_free(commitment_indices);
    enif_free(cell_indices);
    enif_free(cells);
    enif_free(proofs);

    if (ret != C_KZG_OK) {
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_int(env, ret));
    }

    Bytes32 out;
    bytes_from_bls_field(&out, &challenge);

    ErlNifBinary out_bin;
    enif_alloc_binary(BYTES_PER_FIELD_ELEMENT, &out_bin);
    memcpy(out_bin.data, out.bytes, BYTES_PER_FIELD_ELEMENT);
    return enif_make_tuple2(env, enif_make_atom(env, "ok"), enif_make_binary(env, &out_bin));
}

static ErlNifFunc nif_funcs[] = {
    {"load_trusted_setup", 1, nif_load_trusted_setup, 0},
    {"blob_to_kzg_commitment", 1, nif_blob_to_kzg_commitment, 0},
    {"compute_kzg_proof", 2, nif_compute_kzg_proof, 0},
    {"verify_kzg_proof", 4, nif_verify_kzg_proof, 0},
    {"verify_blob_kzg_proof", 3, nif_verify_blob_kzg_proof, 0},
    {"compute_blob_kzg_proof", 2, nif_compute_blob_kzg_proof, 0},
    {"verify_blob_kzg_proof_batch", 3, nif_verify_blob_kzg_proof_batch, 0},
    {"compute_challenge", 2, nif_compute_challenge, 0},
    {"compute_cells_and_kzg_proofs", 1, nif_compute_cells_and_kzg_proofs, 0},
    {"recover_cells_and_kzg_proofs", 2, nif_recover_cells_and_kzg_proofs, 0},
    {"verify_cell_kzg_proof_batch", 4, nif_verify_cell_kzg_proof_batch, 0},
    {"compute_verify_cell_kzg_proof_batch_challenge", 5, nif_compute_verify_cell_kzg_proof_batch_challenge, 0}
};

ERL_NIF_INIT(Elixir.KzgElixir, nif_funcs, NULL, NULL, NULL, NULL)