// NIF bindings for the QNN (Qualcomm AI Engine Direct) runtime.
//
// Resource hierarchy: Graph -> Context -> Runtime, chained with
// enif_keep_resource so destructors run leaf-first. Explicit close/1 exists
// for Runtime/Context because freeing HTP handles can take tens of
// milliseconds and GC destructors run on normal schedulers.

#include <erl_nif.h>

#include <atomic>
#include <cstdio>
#include <ctime>
#include <vector>

#include "nif_utils.hpp"
#include "qnn_api.hpp"
#include "qnn_error.hpp"
#include "qnn_meta.hpp"

using erlang::nif::error;
using erlang::nif::make_binary;
using erlang::nif::ok;

struct RuntimeRes {
    QnnApi api;
    bool closed = false;
};

struct ContextRes {
    Qnn_ContextHandle_t ctx = nullptr;
    NifRes<RuntimeRes> *rt = nullptr;
    std::vector<GraphMeta> graphs;
    bool closed = false;
};

struct GraphRes {
    Qnn_GraphHandle_t graph = nullptr;
    NifRes<ContextRes> *ctx = nullptr;
    const GraphMeta *meta = nullptr;
    ErlNifMutex *mtx = nullptr;
};

template <>
ErlNifResourceType *NifRes<RuntimeRes>::type = nullptr;
template <>
ErlNifResourceType *NifRes<ContextRes>::type = nullptr;
template <>
ErlNifResourceType *NifRes<GraphRes>::type = nullptr;

// ---------------------------------------------------------------------------
// Native stats: leak guards (alive counts must return to baseline after GC)
// and QNN execution accounting. Exposed via the stats/0 NIF.

struct NativeStats {
    std::atomic<int64_t> runtimes_alive{0};
    std::atomic<int64_t> contexts_alive{0};
    std::atomic<int64_t> graphs_alive{0};
    std::atomic<uint64_t> runtimes_created{0};
    std::atomic<uint64_t> contexts_created{0};
    std::atomic<uint64_t> graphs_created{0};
    std::atomic<uint64_t> executes{0};
    std::atomic<uint64_t> execute_errors{0};
    std::atomic<uint64_t> execute_ns_total{0};
    std::atomic<uint64_t> execute_ns_last{0};
};

static NativeStats g_stats;

static uint64_t monotonic_ns() {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return static_cast<uint64_t>(ts.tv_sec) * 1000000000ull + ts.tv_nsec;
}

// ---------------------------------------------------------------------------
// Internal helpers

static void runtime_do_close(RuntimeRes *rt) {
    if (rt && !rt->closed) {
        qnn_api_close(rt->api);
        rt->closed = true;
    }
}

static void context_do_close(ContextRes *ctx) {
    if (ctx && !ctx->closed) {
        RuntimeRes *rt = ctx->rt ? ctx->rt->val : nullptr;
        if (ctx->ctx && rt && !rt->closed && rt->api.fn.contextFree) {
            rt->api.fn.contextFree(ctx->ctx, nullptr);
        }
        ctx->ctx = nullptr;
        ctx->closed = true;
    }
}

static ERL_NIF_TERM quant_term(ErlNifEnv *env, const TensorMeta &t) {
    switch (t.quant) {
        case TensorMeta::Quant::PerTensor:
            return enif_make_tuple3(env,
                                    enif_make_atom(env, "scale_offset"),
                                    enif_make_double(env, t.scale_offset.scale),
                                    enif_make_int(env, t.scale_offset.offset));
        case TensorMeta::Quant::PerAxis: {
            std::vector<ERL_NIF_TERM> pairs;
            pairs.reserve(t.axis_scale_offsets.size());
            for (const auto &so : t.axis_scale_offsets) {
                pairs.push_back(enif_make_tuple2(env,
                                                 enif_make_double(env, so.scale),
                                                 enif_make_int(env, so.offset)));
            }
            return enif_make_tuple3(env,
                                    enif_make_atom(env, "axis_scale_offset"),
                                    enif_make_int(env, t.axis),
                                    enif_make_list_from_array(env, pairs.data(), pairs.size()));
        }
        default:
            return enif_make_atom(env, "none");
    }
}

static ERL_NIF_TERM tensor_meta_term(ErlNifEnv *env, const TensorMeta &t) {
    std::vector<ERL_NIF_TERM> dims;
    dims.reserve(t.dims.size());
    for (uint32_t d : t.dims) dims.push_back(enif_make_uint(env, d));

    ERL_NIF_TERM map = enif_make_new_map(env);
    enif_make_map_put(env, map, enif_make_atom(env, "name"), make_binary(env, t.name), &map);
    enif_make_map_put(env, map, enif_make_atom(env, "id"), enif_make_uint(env, t.id), &map);
    enif_make_map_put(env, map, enif_make_atom(env, "qnn_data_type"),
                      enif_make_uint(env, t.data_type), &map);
    enif_make_map_put(env, map, enif_make_atom(env, "dims"),
                      enif_make_list_from_array(env, dims.data(), dims.size()), &map);
    enif_make_map_put(env, map, enif_make_atom(env, "byte_size"),
                      enif_make_uint64(env, t.byte_size), &map);
    enif_make_map_put(env, map, enif_make_atom(env, "quant"), quant_term(env, t), &map);
    return map;
}

// Load a whole file into memory. Context binaries can be hundreds of MB;
// this runs on a dirty scheduler.
static bool read_file(const std::string &path, std::vector<uint8_t> &out) {
    FILE *f = std::fopen(path.c_str(), "rb");
    if (!f) return false;
    std::fseek(f, 0, SEEK_END);
    long size = std::ftell(f);
    if (size < 0) {
        std::fclose(f);
        return false;
    }
    std::fseek(f, 0, SEEK_SET);
    out.resize(static_cast<size_t>(size));
    size_t read = std::fread(out.data(), 1, out.size(), f);
    std::fclose(f);
    return read == out.size();
}

// ---------------------------------------------------------------------------
// Runtime

static ERL_NIF_TERM runtime_create(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
    (void)argc;
    std::string backend_lib, system_lib;
    int log_level = 0;
    if (!erlang::nif::get_binary_string(env, argv[0], backend_lib) ||
        !erlang::nif::get_binary_string(env, argv[1], system_lib) ||
        !enif_get_int(env, argv[2], &log_level)) {
        return enif_make_badarg(env);
    }

    auto *res = NifRes<RuntimeRes>::allocate(env);
    if (!res) return error(env, "alloc_failed");
    res->val = new RuntimeRes();

    std::string err = qnn_api_open(res->val->api, backend_lib, system_lib, log_level);
    if (!err.empty()) {
        ERL_NIF_TERM reason = make_binary(env, err);
        enif_release_resource(res);
        return error(env, reason);
    }

    g_stats.runtimes_alive++;
    g_stats.runtimes_created++;

    ERL_NIF_TERM term = enif_make_resource(env, res);
    enif_release_resource(res);
    return ok(env, term);
}

static ERL_NIF_TERM runtime_info(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
    (void)argc;
    auto *res = NifRes<RuntimeRes>::get(env, argv[0]);
    if (!res || !res->val) return enif_make_badarg(env);

    ERL_NIF_TERM map = enif_make_new_map(env);
    enif_make_map_put(env, map, enif_make_atom(env, "backend_lib"),
                      make_binary(env, res->val->api.backend_lib_name), &map);
    enif_make_map_put(env, map, enif_make_atom(env, "build_id"),
                      make_binary(env, res->val->api.build_id), &map);
    enif_make_map_put(env, map, enif_make_atom(env, "has_device"),
                      enif_make_atom(env, res->val->api.device ? "true" : "false"), &map);
    enif_make_map_put(env, map, enif_make_atom(env, "closed"),
                      enif_make_atom(env, res->val->closed ? "true" : "false"), &map);
    return map;
}

static ERL_NIF_TERM runtime_close(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
    (void)argc;
    auto *res = NifRes<RuntimeRes>::get(env, argv[0]);
    if (!res || !res->val) return enif_make_badarg(env);
    runtime_do_close(res->val);
    return ok(env);
}

// ---------------------------------------------------------------------------
// Context

static ERL_NIF_TERM context_create(ErlNifEnv *env,
                                   NifRes<RuntimeRes> *rt_res,
                                   const void *buf,
                                   uint64_t size) {
    RuntimeRes *rt = rt_res->val;
    if (rt->closed) return error(env, "runtime_closed");
    if (!rt->api.fn.contextCreateFromBinary) return error(env, "missing_fn");

    auto *res = NifRes<ContextRes>::allocate(env);
    if (!res) return error(env, "alloc_failed");
    res->val = new ContextRes();

    std::string err = qnn_parse_binary_info(rt->api, buf, size, res->val->graphs);
    if (!err.empty()) {
        ERL_NIF_TERM reason = make_binary(env, err);
        enif_release_resource(res);
        return error(env, reason);
    }

    Qnn_ErrorHandle_t qerr = rt->api.fn.contextCreateFromBinary(
        rt->api.backend, rt->api.device, nullptr, buf, size, &res->val->ctx, nullptr);
    if (qerr != QNN_SUCCESS) {
        ERL_NIF_TERM reason = qnn_error_term(env, rt->api, qerr);
        enif_release_resource(res);
        return reason;
    }

    res->val->rt = rt_res;
    enif_keep_resource(rt_res);
    g_stats.contexts_alive++;
    g_stats.contexts_created++;

    ERL_NIF_TERM term = enif_make_resource(env, res);
    enif_release_resource(res);
    return ok(env, term);
}

static ERL_NIF_TERM context_create_from_file(ErlNifEnv *env, int argc,
                                             const ERL_NIF_TERM argv[]) {
    (void)argc;
    auto *rt_res = NifRes<RuntimeRes>::get(env, argv[0]);
    std::string path;
    if (!rt_res || !rt_res->val || !erlang::nif::get_binary_string(env, argv[1], path)) {
        return enif_make_badarg(env);
    }

    std::vector<uint8_t> buf;
    if (!read_file(path, buf)) {
        return error(env, "file_read_failed");
    }
    return context_create(env, rt_res, buf.data(), buf.size());
}

static ERL_NIF_TERM context_create_from_binary(ErlNifEnv *env, int argc,
                                               const ERL_NIF_TERM argv[]) {
    (void)argc;
    auto *rt_res = NifRes<RuntimeRes>::get(env, argv[0]);
    ErlNifBinary bin;
    if (!rt_res || !rt_res->val || !enif_inspect_binary(env, argv[1], &bin)) {
        return enif_make_badarg(env);
    }
    return context_create(env, rt_res, bin.data, bin.size);
}

static ERL_NIF_TERM context_graphs_info(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
    (void)argc;
    auto *res = NifRes<ContextRes>::get(env, argv[0]);
    if (!res || !res->val) return enif_make_badarg(env);

    std::vector<ERL_NIF_TERM> graphs;
    graphs.reserve(res->val->graphs.size());
    for (const auto &g : res->val->graphs) {
        std::vector<ERL_NIF_TERM> ins, outs;
        ins.reserve(g.inputs.size());
        outs.reserve(g.outputs.size());
        for (const auto &t : g.inputs) ins.push_back(tensor_meta_term(env, t));
        for (const auto &t : g.outputs) outs.push_back(tensor_meta_term(env, t));

        ERL_NIF_TERM map = enif_make_new_map(env);
        enif_make_map_put(env, map, enif_make_atom(env, "name"), make_binary(env, g.name), &map);
        enif_make_map_put(env, map, enif_make_atom(env, "inputs"),
                          enif_make_list_from_array(env, ins.data(), ins.size()), &map);
        enif_make_map_put(env, map, enif_make_atom(env, "outputs"),
                          enif_make_list_from_array(env, outs.data(), outs.size()), &map);
        graphs.push_back(map);
    }
    return enif_make_list_from_array(env, graphs.data(), graphs.size());
}

static ERL_NIF_TERM context_close(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
    (void)argc;
    auto *res = NifRes<ContextRes>::get(env, argv[0]);
    if (!res || !res->val) return enif_make_badarg(env);
    context_do_close(res->val);
    return ok(env);
}

// ---------------------------------------------------------------------------
// Graph

static ERL_NIF_TERM graph_retrieve(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
    (void)argc;
    auto *ctx_res = NifRes<ContextRes>::get(env, argv[0]);
    std::string name;
    if (!ctx_res || !ctx_res->val || !erlang::nif::get_binary_string(env, argv[1], name)) {
        return enif_make_badarg(env);
    }
    ContextRes *ctx = ctx_res->val;
    if (ctx->closed) return error(env, "context_closed");
    RuntimeRes *rt = ctx->rt->val;
    if (rt->closed) return error(env, "runtime_closed");
    if (!rt->api.fn.graphRetrieve) return error(env, "missing_fn");

    const GraphMeta *meta = nullptr;
    for (const auto &g : ctx->graphs) {
        if (g.name == name) {
            meta = &g;
            break;
        }
    }
    if (!meta) return error(env, "graph_not_found");

    Qnn_GraphHandle_t handle = nullptr;
    Qnn_ErrorHandle_t qerr = rt->api.fn.graphRetrieve(ctx->ctx, name.c_str(), &handle);
    if (qerr != QNN_SUCCESS) {
        return qnn_error_term(env, rt->api, qerr);
    }

    auto *res = NifRes<GraphRes>::allocate(env);
    if (!res) return error(env, "alloc_failed");
    res->val = new GraphRes();
    res->val->graph = handle;
    res->val->meta = meta;
    res->val->mtx = enif_mutex_create(const_cast<char *>("hexagon_tpu_graph"));
    res->val->ctx = ctx_res;
    enif_keep_resource(ctx_res);
    g_stats.graphs_alive++;
    g_stats.graphs_created++;

    ERL_NIF_TERM term = enif_make_resource(env, res);
    enif_release_resource(res);
    return ok(env, term);
}

// Build a v1 execute tensor descriptor from cached metadata. `dims` must
// outlive the QNN call (it points into the context's GraphMeta).
static Qnn_Tensor_t make_exec_tensor(const TensorMeta &meta, bool is_input, void *data,
                                     uint32_t data_size) {
    Qnn_Tensor_t t = QNN_TENSOR_INIT;
    t.version = QNN_TENSOR_VERSION_1;
    t.v1.id = meta.id;
    t.v1.name = meta.name.c_str();
    t.v1.type = is_input ? QNN_TENSOR_TYPE_APP_WRITE : QNN_TENSOR_TYPE_APP_READ;
    t.v1.dataFormat = QNN_TENSOR_DATA_FORMAT_DENSE;
    t.v1.dataType = meta.data_type;
    if (meta.quant == TensorMeta::Quant::PerTensor) {
        t.v1.quantizeParams.encodingDefinition = QNN_DEFINITION_DEFINED;
        t.v1.quantizeParams.quantizationEncoding = QNN_QUANTIZATION_ENCODING_SCALE_OFFSET;
        t.v1.quantizeParams.scaleOffsetEncoding.scale = meta.scale_offset.scale;
        t.v1.quantizeParams.scaleOffsetEncoding.offset = meta.scale_offset.offset;
    }
    t.v1.rank = static_cast<uint32_t>(meta.dims.size());
    t.v1.dimensions = const_cast<uint32_t *>(meta.dims.data());
    t.v1.memType = QNN_TENSORMEMTYPE_RAW;
    t.v1.clientBuf.data = data;
    t.v1.clientBuf.dataSize = data_size;
    return t;
}

static ERL_NIF_TERM graph_execute(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
    (void)argc;
    auto *graph_res = NifRes<GraphRes>::get(env, argv[0]);
    if (!graph_res || !graph_res->val) return enif_make_badarg(env);
    GraphRes *g = graph_res->val;
    ContextRes *ctx = g->ctx->val;
    RuntimeRes *rt = ctx->rt->val;
    if (ctx->closed) return error(env, "context_closed");
    if (rt->closed) return error(env, "runtime_closed");
    if (!rt->api.fn.graphExecute) return error(env, "missing_fn");

    const auto &in_meta = g->meta->inputs;
    const auto &out_meta = g->meta->outputs;

    // Inputs: zero-copy views into the caller's binaries. The terms stay
    // live for the duration of this (synchronous, dirty) call.
    unsigned list_len = 0;
    if (!enif_get_list_length(env, argv[1], &list_len) || list_len != in_meta.size()) {
        return error(env, "input_count_mismatch");
    }

    std::vector<Qnn_Tensor_t> inputs;
    inputs.reserve(in_meta.size());
    ERL_NIF_TERM list = argv[1], head, tail;
    for (size_t i = 0; i < in_meta.size(); i++) {
        enif_get_list_cell(env, list, &head, &tail);
        list = tail;
        ErlNifBinary bin;
        if (!enif_inspect_binary(env, head, &bin)) {
            return enif_make_badarg(env);
        }
        if (in_meta[i].byte_size != 0 && bin.size != in_meta[i].byte_size) {
            return error(env,
                         enif_make_tuple3(env, enif_make_atom(env, "input_size_mismatch"),
                                          enif_make_uint64(env, i),
                                          enif_make_uint64(env, in_meta[i].byte_size)));
        }
        inputs.push_back(make_exec_tensor(in_meta[i], true, bin.data,
                                          static_cast<uint32_t>(bin.size)));
    }

    // Outputs: owned binaries handed to QNN to fill, then transferred to
    // the returned terms without copying.
    std::vector<ErlNifBinary> out_bins(out_meta.size());
    std::vector<Qnn_Tensor_t> outputs;
    outputs.reserve(out_meta.size());
    for (size_t i = 0; i < out_meta.size(); i++) {
        if (!enif_alloc_binary(out_meta[i].byte_size, &out_bins[i])) {
            for (size_t j = 0; j < i; j++) enif_release_binary(&out_bins[j]);
            return error(env, "alloc_failed");
        }
        outputs.push_back(make_exec_tensor(out_meta[i], false, out_bins[i].data,
                                           static_cast<uint32_t>(out_bins[i].size)));
    }

    enif_mutex_lock(g->mtx);
    uint64_t t0 = monotonic_ns();
    Qnn_ErrorHandle_t qerr = rt->api.fn.graphExecute(
        g->graph, inputs.data(), static_cast<uint32_t>(inputs.size()), outputs.data(),
        static_cast<uint32_t>(outputs.size()), nullptr, nullptr);
    uint64_t elapsed = monotonic_ns() - t0;
    enif_mutex_unlock(g->mtx);

    g_stats.executes++;
    g_stats.execute_ns_total += elapsed;
    g_stats.execute_ns_last = elapsed;

    if (qerr != QNN_SUCCESS) {
        g_stats.execute_errors++;
        for (auto &b : out_bins) enif_release_binary(&b);
        return qnn_error_term(env, rt->api, qerr);
    }

    std::vector<ERL_NIF_TERM> out_terms;
    out_terms.reserve(out_bins.size());
    for (auto &b : out_bins) {
        out_terms.push_back(enif_make_binary(env, &b));
    }
    return ok(env, enif_make_list_from_array(env, out_terms.data(), out_terms.size()));
}

// ---------------------------------------------------------------------------
// Resource destructors and registration

static void runtime_dtor(ErlNifEnv *, void *obj) {
    auto *res = static_cast<NifRes<RuntimeRes> *>(obj);
    if (res->val) {
        runtime_do_close(res->val);
        delete res->val;
        res->val = nullptr;
        g_stats.runtimes_alive--;
    }
}

static void context_dtor(ErlNifEnv *, void *obj) {
    auto *res = static_cast<NifRes<ContextRes> *>(obj);
    if (res->val) {
        context_do_close(res->val);
        if (res->val->rt) enif_release_resource(res->val->rt);
        delete res->val;
        res->val = nullptr;
        g_stats.contexts_alive--;
    }
}

static void graph_dtor(ErlNifEnv *, void *obj) {
    auto *res = static_cast<NifRes<GraphRes> *>(obj);
    if (res->val) {
        // Graph handles are owned by the context; nothing to free in QNN.
        if (res->val->mtx) enif_mutex_destroy(res->val->mtx);
        if (res->val->ctx) enif_release_resource(res->val->ctx);
        delete res->val;
        res->val = nullptr;
        g_stats.graphs_alive--;
    }
}

// stats/0 — native resource + execution accounting (leak guard).
static ERL_NIF_TERM stats(ErlNifEnv *env, int, const ERL_NIF_TERM[]) {
    ERL_NIF_TERM map = enif_make_new_map(env);
    auto put_i64 = [&](const char *k, int64_t v) {
        enif_make_map_put(env, map, enif_make_atom(env, k), enif_make_int64(env, v), &map);
    };
    auto put_u64 = [&](const char *k, uint64_t v) {
        enif_make_map_put(env, map, enif_make_atom(env, k), enif_make_uint64(env, v), &map);
    };
    put_i64("runtimes_alive", g_stats.runtimes_alive.load());
    put_i64("contexts_alive", g_stats.contexts_alive.load());
    put_i64("graphs_alive", g_stats.graphs_alive.load());
    put_u64("runtimes_created", g_stats.runtimes_created.load());
    put_u64("contexts_created", g_stats.contexts_created.load());
    put_u64("graphs_created", g_stats.graphs_created.load());
    put_u64("executes", g_stats.executes.load());
    put_u64("execute_errors", g_stats.execute_errors.load());
    put_u64("execute_ns_total", g_stats.execute_ns_total.load());
    put_u64("execute_ns_last", g_stats.execute_ns_last.load());
    return map;
}

static int on_load(ErlNifEnv *env, void **, ERL_NIF_TERM) {
    NifRes<RuntimeRes>::type =
        enif_open_resource_type(env, nullptr, "hexagon_tpu_runtime", runtime_dtor,
                                (ErlNifResourceFlags)(ERL_NIF_RT_CREATE | ERL_NIF_RT_TAKEOVER),
                                nullptr);
    NifRes<ContextRes>::type =
        enif_open_resource_type(env, nullptr, "hexagon_tpu_context", context_dtor,
                                (ErlNifResourceFlags)(ERL_NIF_RT_CREATE | ERL_NIF_RT_TAKEOVER),
                                nullptr);
    NifRes<GraphRes>::type =
        enif_open_resource_type(env, nullptr, "hexagon_tpu_graph", graph_dtor,
                                (ErlNifResourceFlags)(ERL_NIF_RT_CREATE | ERL_NIF_RT_TAKEOVER),
                                nullptr);
    if (!NifRes<RuntimeRes>::type || !NifRes<ContextRes>::type || !NifRes<GraphRes>::type) {
        return -1;
    }
    return 0;
}

static int on_upgrade(ErlNifEnv *env, void **priv, void **, ERL_NIF_TERM info) {
    return on_load(env, priv, info);
}

static ErlNifFunc nif_functions[] = {
    {"runtime_create", 3, runtime_create, ERL_NIF_DIRTY_JOB_IO_BOUND},
    {"runtime_info", 1, runtime_info, 0},
    {"runtime_close", 1, runtime_close, ERL_NIF_DIRTY_JOB_IO_BOUND},
    {"context_create_from_file", 2, context_create_from_file, ERL_NIF_DIRTY_JOB_CPU_BOUND},
    {"context_create_from_binary", 2, context_create_from_binary, ERL_NIF_DIRTY_JOB_CPU_BOUND},
    {"context_graphs_info", 1, context_graphs_info, 0},
    {"context_close", 1, context_close, ERL_NIF_DIRTY_JOB_IO_BOUND},
    {"graph_retrieve", 2, graph_retrieve, ERL_NIF_DIRTY_JOB_CPU_BOUND},
    {"graph_execute", 2, graph_execute, ERL_NIF_DIRTY_JOB_CPU_BOUND},
    {"stats", 0, stats, 0},
};

ERL_NIF_INIT(Elixir.HexagonTpu.Nif, nif_functions, on_load, nullptr, on_upgrade, nullptr)