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

#include "ckzg.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 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}
};

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