#pragma clang diagnostic ignored "-Wgnu-zero-variadic-macro-arguments"
#pragma clang diagnostic ignored "-Wunused-function"
#pragma clang diagnostic ignored "-Wunused-variable"
#pragma clang diagnostic ignored "-Wunused-but-set-variable"

#include <HAP_farf.h>
#include <HAP_perf.h>
#include <AEEStdErr.h>
#include <dspqueue.h>
#include <HAP_compute_res.h>
#include <HAP_etm_config.h>
#include <HAP_mem.h>
#include <HAP_power.h>
#include <HAP_ps.h>
#include <HAP_dcvs.h>
#include <qurt.h>
#include <qurt_thread.h>
#include <qurt_memory.h>
#include <remote.h>
#include <string.h>
#include <stdatomic.h>

#include "hex-utils.h"
#include "hex-dma.h"
#include "hmx-queue.h"

#define GGML_COMMON_DECL_C
#include "ggml-common.h"
#include "hex-bitmap.h"
#include "htp-ctx.h"
#include "htp-ops.h"
#include "htp-tensor.h"
#include "htp_iface.h"
#include "work-queue.h"
#include "hex-profile.h"
#include "allreduce-ops.h"

#define HMX_QUEUE_CAPACITY     16
#define HMX_QUEUE_STACK_SIZE   16384
#define WORK_QUEUE_CAPACITY    16
#define WORK_QUEUE_STACK_SIZE  16384
#define MAIN_THREAD_STACK_SIZE 32768

_Static_assert(WORK_QUEUE_MAX_N_THREADS >= HTP_MAX_NTHREADS,
               "work-queue thread cap must be >= HTP_MAX_NTHREADS");

struct htp_handle {
    struct htp_context * ctx;
};

static inline void * htp_mmap(uint32_t fd, uint32_t size) {
    void * va = (void *)-1;
    for (int retry = 0; retry < 2; retry++) {
#if __HVX_ARCH__ > 73
        va = HAP_mmap2(NULL, size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0);
#else
        if (size > HTP_MMAP_MAX_VMEM) {
            FARF(ERROR, "mmap failed : size %u exceeds 2GB limit for HAP_mmap", (uint32_t) size);
            abort();
        }
        va = HAP_mmap(NULL, size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0);
#endif
        if (va != (void *)-1 && va != NULL) {
            return va;
        }
        if (retry == 0) {
            FARF(HIGH, "mmap failed first try (va %p fd %u size %u), retrying...", va, fd, size);
        }
    }
    return NULL;
}

static inline void htp_munmap(void * va, uint32_t size) {
#if __HVX_ARCH__ > 73
    HAP_munmap2(va, size);
#else
    HAP_munmap(va, size);
#endif
}

AEEResult htp_iface_open(const char * uri, remote_handle64 * handle) {
    (void) uri;
    struct htp_handle * h = calloc(1, sizeof(*h));
    if (h == NULL) {
        return AEE_ENOMEMORY;
    }

    *handle = (remote_handle64) h;
    return AEE_SUCCESS;
}

AEEResult htp_iface_etm(remote_handle64 handle, uint32_t enable) {
    struct htp_handle * h = (struct htp_handle *) handle;
    if (!h) {
        return AEE_EBADPARM;
    }

    int err = enable ? HAP_user_etm_enable() : HAP_user_etm_disable();
    if (err) {
        if (err == AEE_EVERSIONNOTSUPPORT) {
            FARF(ERROR, "API HAP_user_etm_enable/disable is not supported\n");
        } else {
            FARF(ERROR, "Error executing HAP_user_etm_enable/disable with error code : 0x%x\n", err);
        }
    }
    return err;
}

AEEResult htp_iface_profiler(remote_handle64 handle, uint32_t mode, const htp_iface_pmu_conf* pmu_conf) {
    struct htp_handle * h = (struct htp_handle *) handle;
    if (!h || !h->ctx) {
        return AEE_EBADPARM;
    }
    struct htp_context * ctx = h->ctx;

    if (mode == HTP_PROF_PMU) {
        const uint32_t* events = pmu_conf->events;

        // Pack 4 event IDs (low 8 bits) into each 32-bit config register
        uint32_t evtcfg = 0, evtcfg1 = 0, cfg = 0, i = 0;
        for (; i < HEX_NUM_PMU_COUNTERS/2; i++) {
            evtcfg  |= ((events[i + 0] & 0xFF) << (i * 8));
            evtcfg1 |= ((events[i + 4] & 0xFF) << (i * 8));
        }

        // For events >255 pack high 2 bits of all 8 event IDs into cfg register
        // 2 bits per counter: bits [1:0] for counter 0, [3:2] for counter 1, etc.
        for (i = 0; i < HEX_NUM_PMU_COUNTERS; i++) {
            cfg |= (((events[i] >> 8) & 3) << (i * 2));
        }

        FARF(ALWAYS, "Configuring PMU registers: evtcfg = 0x%x, evtcfg1 = 0x%x, pmucfg = 0x%x", evtcfg, evtcfg1, cfg);

        // Configure PMU registers
        qurt_pmu_set(QURT_PMUCFG,     cfg);
        qurt_pmu_set(QURT_PMUEVTCFG,  evtcfg);
        qurt_pmu_set(QURT_PMUEVTCFG1, evtcfg1);
        qurt_pmu_enable(1);
    }

    ctx->profiler = mode;

    return AEE_SUCCESS;
}

AEEResult htp_iface_close(remote_handle64 handle) {
    struct htp_handle * h = (struct htp_handle *) handle;
    if (!h) {
        return AEE_EBADPARM;
    }

    struct htp_context * ctx = h->ctx;
    if (ctx) {
        if (ctx->dsp_queue) {
            FARF(ERROR, "Closing handle with queue still open");
            return AEE_EITEMBUSY;
        }

        // release the mmaps (if any)
        for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
            if (ctx->mmap[i].size) {
                htp_munmap((void *) ctx->mmap[i].base, ctx->mmap[i].size);
                ctx->mmap[i].size = 0;
                ctx->mmap[i].base = NULL;
                ctx->mmap[i].fd   = -1;
            }
        }

        if (ctx->profiler) {
            qurt_pmu_enable(1);
        }

        if (ctx->etm) {
            HAP_user_etm_disable();
        }

        // Free the unified block (ctx is the base address of the block)
        free(ctx);
        h->ctx = NULL;
    }

    free(h);
    return AEE_SUCCESS;
}

AEEResult htp_iface_mmap(remote_handle64 handle, uint32_t fd, uint32_t size) {
    struct htp_handle * h = (struct htp_handle *) handle;
    if (!h || !h->ctx) {
        return AEE_EBADPARM;
    }
    struct htp_context * ctx = h->ctx;

    // See if we already have this mapping
    for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
        struct htp_mmap *m = &ctx->mmap[i];
        if (m->fd == fd) {
            return AEE_SUCCESS;
        }
    }

    // Add new mapping
    for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
        struct htp_mmap *m = &ctx->mmap[i];
        if (!m->size) {
            FARF(HIGH, "mmap : fd %u size %u", fd, size);
            void *va = htp_mmap(fd, size);
            if (va == NULL) {
                FARF(ERROR, "mmap failed : fd %u size %u", fd, (uint32_t) size);
                return AEE_EFAILED;
            }

            m->base   = (uint64_t) va;
            m->fd     = fd;
            m->size   = size;

            return AEE_SUCCESS;
        }
    }

    return AEE_ENOMEMORY;
}

AEEResult htp_iface_munmap(remote_handle64 handle, uint32 fd) {
    struct htp_handle * h = (struct htp_handle *) handle;
    if (!h || !h->ctx) {
        return AEE_EBADPARM;
    }
    struct htp_context * ctx = h->ctx;

    for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
        struct htp_mmap *m = &ctx->mmap[i];
        if (fd < 0 || m->fd == fd) {
            FARF(HIGH, "unmmap : base %p fd %u size %u", (void*) m->base, m->fd, (uint32_t) m->size);
            htp_munmap((void *) m->base, m->size);
            m->size   = 0;
            m->base   = NULL;
            m->fd     = -1;
        }
    }

    return AEE_SUCCESS;
}

static void vtcm_acquire(struct htp_context * ctx) {
    if (!ctx->vtcm_valid) {
        int err = HAP_compute_res_acquire_cached(ctx->vtcm_rctx, 10000000u);
        if (err != 0) {
            FARF(ERROR, "ggml-hex: failed to acquire VTCM: 0x%08x", (unsigned)err);
            abort();
        }

        ctx->vtcm_needs_release = false;
        ctx->vtcm_valid = true;

        // Drop the priority to make sure we get the release callback from other GGML-HTP and QNN-HTP sessions
        HAP_compute_res_update_priority(ctx->vtcm_rctx, ctx->thread_prio + 10);
    }
}

static void vtcm_release(struct htp_context * ctx) {
    if (ctx->vtcm_valid) {
        ctx->vtcm_valid         = false;
        ctx->vtcm_needs_release = false;
        HAP_compute_res_release_cached(ctx->vtcm_rctx);
    }
}

static int vtcm_release_callback(unsigned int rctx, void * state) {
    struct htp_context * ctx = (struct htp_context *) state;
    ctx->vtcm_needs_release = true;
    return 0;
}

static int vtcm_alloc(struct htp_context * ctx) {
    unsigned int vtcm_size = 8 * 1024 * 1024;  // 8MB default
    HAP_compute_res_query_VTCM(0, &vtcm_size, NULL, NULL, NULL);

    compute_res_attr_t attr;
    HAP_compute_res_attr_init(&attr);
    HAP_compute_res_attr_set_serialize(&attr, 0);
    HAP_compute_res_attr_set_cache_mode(&attr, 1);
    HAP_compute_res_attr_set_vtcm_param_v2(&attr, vtcm_size, vtcm_size, vtcm_size); // single page
    HAP_compute_res_attr_set_release_callback(&attr, vtcm_release_callback, (void *) ctx);
    HAP_compute_res_attr_set_hmx_param(&attr, 1);

    // Allocate VTCM for scratch pads
    uint32_t rctx = HAP_compute_res_acquire(&attr, 1000000 /* timeout */);
    if (!rctx) {
        FARF(ERROR, "failed to allocate %zu bytes VTCM\n", ctx->vtcm_size);
        return AEE_ENOMEMORY;
    }

    void * vtcm_ptr;
    if (HAP_compute_res_attr_get_vtcm_ptr_v2(&attr, &vtcm_ptr, &vtcm_size) != 0) {
        HAP_compute_res_release(rctx);
        FARF(ERROR, "failed to allocate %zu bytes VTCM (new)\n", ctx->vtcm_size);
        return AEE_ENOMEMORY;
    }

    ctx->vtcm_base          = (uint8_t *) vtcm_ptr;
    ctx->vtcm_size          = vtcm_size;
    ctx->vtcm_rctx          = rctx;
    ctx->vtcm_valid         = false;
    ctx->vtcm_needs_release = false;

    return 0;
}

static void vtcm_free(struct htp_context * ctx) {
    if (ctx->vtcm_rctx) {
        HAP_compute_res_release(ctx->vtcm_rctx);
        ctx->vtcm_base = 0;
        ctx->vtcm_rctx = 0;
    }
}

static void htp_main_thread(void * context);
static void htp_packet_callback(dspqueue_t queue, int error, void * context);
static void htp_error_callback(dspqueue_t queue, int error, void * context);

AEEResult htp_iface_start(remote_handle64 handle, uint32_t sess_id, uint64_t dsp_queue_id, uint32_t n_hvx, uint32_t n_hmx, uint64_t max_vmem) {
    struct htp_handle * h = (struct htp_handle *) handle;
    if (!h) {
        return AEE_EBADPARM;
    }

    if (h->ctx) {
        FARF(ERROR, "Queue already open");
        return AEE_EITEMBUSY;
    }

    // Cache the original FastRPC thread priority, then calculate compute priority
    int fastrpc_tid  = qurt_thread_get_id();
    int fastrpc_prio = qurt_thread_get_priority(fastrpc_tid);
    int main_prio    = fastrpc_prio - 10;
    if (main_prio < 1) main_prio = 1;

    dspqueue_t dsp_queue = NULL;
    bool use_callbacks = false;

    // Import queue with NULL callbacks to avoid starting dspueue internal threads
    int err = dspqueue_import(dsp_queue_id, NULL, NULL, (void *) h, &dsp_queue);
    if (err == AEE_EBADPARM) {
        // Fallback for devices that don't support NULL callbacks
        FARF(HIGH, "dspqueue import with NULL callbacks failed, trying with callbacks");
        use_callbacks = true;
        err = dspqueue_import(dsp_queue_id, htp_packet_callback, htp_error_callback, (void *) h, &dsp_queue);
    }

    if (err) {
        FARF(ERROR, "Queue import failed with 0x%08x", (unsigned) err);
        return err;
    }

    qurt_sysenv_max_hthreads_t hw_threads;
    qurt_sysenv_get_max_hw_threads(&hw_threads);
    uint32_t hw_nhvx = (qurt_hvx_get_units() >> 8) & 0xFF;

    if (n_hvx == 0) {
        n_hvx = hw_nhvx;
    }
    if (n_hvx > hw_threads.max_hthreads) {
        n_hvx = hw_threads.max_hthreads;
    }
    if (n_hvx > HTP_MAX_NTHREADS) {
        n_hvx = HTP_MAX_NTHREADS;
    }

    // layout segments of our contiguous block

    // 1. htp_context : sits at the base (block is 4K-aligned via memalign below)
    size_t offset = sizeof(struct htp_context);

    // 2. main_stack
    size_t offset_main_stack = 0;
    size_t size_main_stack   = 0;
    if (!use_callbacks) {
        offset_main_stack = hex_align_up(offset, 4096);
        size_main_stack   = MAIN_THREAD_STACK_SIZE;
        offset = offset_main_stack + size_main_stack;
    }

    // 3. work_queue
    size_t wq_align  = work_queue_alignof();
    size_t offset_wq = hex_align_up(offset, wq_align);
    size_t size_wq   = work_queue_sizeof(n_hvx, WORK_QUEUE_CAPACITY, WORK_QUEUE_STACK_SIZE);
    offset = offset_wq + size_wq;

    // 4. dma_queue
    size_t dma_align = dma_queue_alignof();
    size_t offset_dma = hex_align_up(offset, dma_align);
    size_t size_dma = 0;
    for (uint32_t i = 0; i < n_hvx; i++) {
        size_dma  = hex_align_up(size_dma, dma_queue_alignof());
        size_dma += dma_queue_sizeof(256);
        size_dma  = hex_align_up(size_dma, dma_queue_alignof());
        size_dma += dma_queue_alias_sizeof();
    }
    offset = offset_dma + size_dma;

    // 5. hmx_queue
    size_t offset_hmx = 0;
    size_t size_hmx = 0;
    if (n_hmx) {
        size_t hmx_align = hmx_queue_alignof();
        offset_hmx = hex_align_up(offset, hmx_align);
        size_hmx   = hmx_queue_sizeof(HMX_QUEUE_CAPACITY, HMX_QUEUE_STACK_SIZE);
        offset = offset_hmx + size_hmx;
    }

    size_t footprint = hex_align_up(offset, 128);

    void * block = memalign(4096, footprint);
    if (!block) {
        FARF(ERROR, "Unable to allocate unified block of size %zu\n", footprint);
        dspqueue_close(dsp_queue);
        return AEE_ENOMEMORY;
    }
    memset(block, 0, footprint);

    h->ctx = (struct htp_context *) block;
    struct htp_context * ctx = h->ctx;
    ctx->footprint = footprint;

    ctx->thread_id   = fastrpc_tid;
    ctx->thread_prio = main_prio;
    ctx->max_vmem    = max_vmem;
    ctx->dsp_queue   = dsp_queue;

    err = vtcm_alloc(ctx);
    if (err != AEE_SUCCESS) {
        FARF(ERROR, "Unable to allocate VTCM");
        htp_iface_stop(handle);
        return AEE_ENOMEMORY;
    }

    HAP_setFARFRuntimeLoggingParams(0xffff, NULL, 0);

    // Set client class
    {
        HAP_power_request_t request;
        memset(&request, 0, sizeof(HAP_power_request_t));
        request.type    = HAP_power_set_apptype;
        request.apptype = HAP_POWER_COMPUTE_CLIENT_CLASS;

        if ((err = HAP_power_set((void *) ctx, &request)) != 0) {
            htp_iface_stop(handle);
            return err;
        }
    }

    // DCVS setup
    {
        HAP_power_request_t request;
        memset(&request, 0, sizeof(request));

        request.type                              = HAP_power_set_DCVS_v3;
        request.dcvs_v3.set_dcvs_enable           = TRUE;
        request.dcvs_v3.dcvs_enable               = FALSE;
        request.dcvs_v3.set_bus_params            = TRUE;
        request.dcvs_v3.bus_params.min_corner     = HAP_DCVS_VCORNER_MAX;
        request.dcvs_v3.bus_params.max_corner     = HAP_DCVS_VCORNER_MAX;
        request.dcvs_v3.bus_params.target_corner  = HAP_DCVS_VCORNER_MAX;
        request.dcvs_v3.set_core_params           = TRUE;
        request.dcvs_v3.core_params.min_corner    = HAP_DCVS_VCORNER_MAX;
        request.dcvs_v3.core_params.max_corner    = HAP_DCVS_VCORNER_MAX;
        request.dcvs_v3.core_params.target_corner = HAP_DCVS_VCORNER_MAX;
        request.dcvs_v3.set_sleep_disable         = TRUE;
        request.dcvs_v3.sleep_disable             = TRUE;

#if (__HEXAGON_ARCH__ >= 79)
        HAP_set_dcvs_v3_protected_bus_corners(&request, 1);
#endif
        if ((err = HAP_power_set((void *) ctx, &request)) != 0) {
            htp_iface_stop(handle);
            return err;
        }

        memset(&request, 0, sizeof(request));
        request.type         = HAP_power_set_HVX;
        request.hvx.power_up = TRUE;
        if ((err = HAP_power_set((void *) ctx, &request)) != 0) {
            htp_iface_stop(handle);
            return err;
        }
    }

#if __HVX_ARCH__ >= 75
    {
        // Power on HMX and set HMX clock
        HAP_power_request_t request;
        memset(&request, 0, sizeof(HAP_power_request_t));
        request.type = HAP_power_set_HMX_v2;
        request.hmx_v2.set_power     = TRUE;
        request.hmx_v2.power_up      = TRUE;
        request.hmx_v2.set_clock     = TRUE;
        request.hmx_v2.target_corner = HAP_DCVS_EXP_VCORNER_MAX;
        request.hmx_v2.min_corner    = HAP_DCVS_EXP_VCORNER_MAX;
        request.hmx_v2.max_corner    = HAP_DCVS_EXP_VCORNER_MAX;
        request.hmx_v2.perf_mode     = HAP_CLK_PERF_HIGH;
        FARF(ALWAYS, "Setting HMX clock\n");
        err = HAP_power_set((void *) ctx, &request);
        if (err != AEE_SUCCESS) {
            FARF(ERROR, "ggml-hex: error setting HMX clock.");
            htp_iface_stop(handle);
            return err;
        }
    }
#else
    {
        // Power on HMX
        HAP_power_request_t request;
        memset(&request, 0, sizeof(HAP_power_request_t));
        request.type         = HAP_power_set_HMX;
        request.hmx.power_up = TRUE;
        FARF(ALWAYS, "Powering HMX on\n");
        err = HAP_power_set((void *) ctx, &request);
        if (err != AEE_SUCCESS) {
            FARF(ERROR, "ggml-hex: error powering on HMX.");
            htp_iface_stop(handle);
            return err;
        }
    }
#endif

    ctx->hmx_enabled = n_hmx;
    ctx->hmx_queue   = NULL;
    if (n_hmx) {
        void * hmx_ptr = (void *) ((uintptr_t) block + offset_hmx);
        ctx->hmx_queue = hmx_queue_init(hmx_ptr, HMX_QUEUE_CAPACITY, HMX_QUEUE_STACK_SIZE, ctx->vtcm_rctx, &ctx->trace[HTP_MAX_NTHREADS]);
    }
    FARF(HIGH, "HMX %s (n_hmx=%d)", ctx->hmx_enabled ? "enabled" : "disabled", n_hmx);

    ctx->n_threads = n_hvx;
    ctx->n_threads_div = init_fastdiv_values(ctx->n_threads);

    // Initialize DMA queues
    uint8_t * dma_ptr_curr = (uint8_t *) ((uintptr_t) block + offset_dma);
    size_t size_dma_q = dma_queue_sizeof(256);
    size_t size_dma_alias = dma_queue_alias_sizeof();

    for (int i = 0; i < ctx->n_threads; i++) {
        dma_ptr_curr = (uint8_t *) hex_align_up((uintptr_t) dma_ptr_curr, dma_queue_alignof());
        ctx->dma_cached[i] = dma_queue_init(dma_ptr_curr, 256, (uintptr_t) ctx->vtcm_base, ctx->vtcm_size, &ctx->trace[i]);
        dma_ptr_curr += size_dma_q;

        dma_ptr_curr = (uint8_t *) hex_align_up((uintptr_t) dma_ptr_curr, dma_queue_alignof());
        ctx->dma[i] = dma_queue_alias_init(dma_ptr_curr, ctx->dma_cached[i], 1);
        dma_ptr_curr += size_dma_alias;
    }

    ctx->ddr_spad_size = 512 * 1024; // 512 KB
    ctx->ddr_spad_base = memalign(128, ctx->ddr_spad_size);

    void * wq_ptr = (void *) ((uintptr_t) block + offset_wq);
    ctx->work_queue = work_queue_init(wq_ptr, n_hvx, WORK_QUEUE_CAPACITY, WORK_QUEUE_STACK_SIZE);

    ctx->main_stack = NULL;
    ctx->main_thread = 0;
    atomic_store(&ctx->killed, false);

    if (!use_callbacks) {
        // Start main compute thread
        ctx->main_stack = (void *) ((uintptr_t) block + offset_main_stack);

        qurt_thread_attr_t attr;
        qurt_thread_attr_init(&attr);
        qurt_thread_attr_set_stack_addr(&attr, ctx->main_stack);
        qurt_thread_attr_set_stack_size(&attr, size_main_stack);
        qurt_thread_attr_set_priority(&attr, main_prio);
        qurt_thread_attr_set_name(&attr, "htp-main");

        int err_thread = qurt_thread_create(&ctx->main_thread, &attr, htp_main_thread, ctx);
        if (err_thread) {
            FARF(ERROR, "Unable to create htp main thread: %d", err_thread);
            htp_iface_stop(handle);
            return AEE_ENOMEMORY;
        }
    }

    FARF(HIGH, "session %u started: n-hvx %u vtcm-size %zu vtcm-rctx %u n-threads %u thread-id %d thread-prio %d \n",
         sess_id, hw_nhvx, ctx->vtcm_size, ctx->vtcm_rctx, ctx->n_threads, ctx->thread_id, ctx->thread_prio);

    return AEE_SUCCESS;
}

AEEResult htp_iface_stop(remote_handle64 handle) {
    struct htp_handle * h = (struct htp_handle *) handle;
    if (!h || !h->ctx) {
        return AEE_EBADPARM;
    }
    struct htp_context * ctx = h->ctx;

    if (ctx->main_thread) {
        atomic_store(&ctx->killed, true);
        int status;
        (void) qurt_thread_join(ctx->main_thread, &status);
        ctx->main_thread = 0;
    }

    int err = dspqueue_close(ctx->dsp_queue); ctx->dsp_queue = NULL;
    if (err != 0) {
        FARF(ERROR, "Queue close failed with 0x%08x", (unsigned) err);
        return err;
    }

    work_queue_free(ctx->work_queue);

    for (int i = 0; i < ctx->n_threads; i++) {
        dma_queue_alias_free(ctx->dma[i]);
        dma_queue_free(ctx->dma_cached[i]);
    }

    if (ctx->hmx_queue) {
        hmx_queue_free(ctx->hmx_queue);
        ctx->hmx_queue = NULL;
    }
    ctx->hmx_enabled = false;

    vtcm_free(ctx);

    if (ctx->ddr_spad_base) {
        free(ctx->ddr_spad_base);
        ctx->ddr_spad_base = NULL;
        ctx->ddr_spad_size = 0;
    }

    free(ctx);
    h->ctx = NULL;

    return AEE_SUCCESS;
}

AEEResult htp_iface_hwinfo(remote_handle64 handle, uint32_t * n_threads, uint32_t * n_hvx, uint32_t * n_hmx, uint64_t * vtcm_size) {
    (void)handle;
    if (!n_threads || !n_hvx || !n_hmx || !vtcm_size) {
        return AEE_EBADPARM;
    }

    qurt_sysenv_max_hthreads_t hw_threads;
    qurt_sysenv_get_max_hw_threads(&hw_threads);
    uint32_t hw_nhvx = (qurt_hvx_get_units() >> 8) & 0xFF;

    uint32_t n_hvx_val = hw_nhvx;
    if (n_hvx_val > hw_threads.max_hthreads) {
        n_hvx_val = hw_threads.max_hthreads;
    }
    if (n_hvx_val > HTP_MAX_NTHREADS) {
        n_hvx_val = HTP_MAX_NTHREADS;
    }

    // for now we force n_threads == n_hvx
    *n_threads = n_hvx_val;
    *n_hvx     = n_hvx_val;
    *n_hmx     = 1;

    uint32_t vtcm_sz = 8 * 1024 * 1024; // 8MB default fallback
    HAP_compute_res_query_VTCM(0, (unsigned int *)&vtcm_sz, NULL, NULL, NULL);
    *vtcm_size = vtcm_sz;

    return AEE_SUCCESS;
}

static void htp_error_callback(dspqueue_t queue, int error, void * context) {
    // No errors expected on the DSP.
    FARF(ERROR, "Error callback: 0x%08x", (unsigned) error);
}

struct profile_data {
    uint64_t usecs;
    uint64_t cycles_start;
    uint64_t cycles_stop;
    uint32_t pmu_counters[HEX_NUM_PMU_COUNTERS];
};

static inline void profile_start(uint32_t mode, struct profile_data * d) {
    switch (mode) {
        case HTP_PROF_PMU:
            hex_get_pmu(d->pmu_counters);
            // fallthrough
        case HTP_PROF_BASIC:
        case HTP_PROF_TRACE:
            d->usecs  = HAP_perf_get_qtimer_count();
            d->cycles_start = hex_get_cycles();
            break;
        default:
            break;
    }
}

static inline void profile_stop(uint32_t mode, struct profile_data * d) {
    uint32_t pmu_counters[HEX_NUM_PMU_COUNTERS];
    switch (mode) {
        case HTP_PROF_PMU:
            hex_get_pmu(pmu_counters);
            for (int i = 0; i < HEX_NUM_PMU_COUNTERS; i++) {
                d->pmu_counters[i] = pmu_counters[i] - d->pmu_counters[i];
            }
            // fallthrough
        case HTP_PROF_BASIC:
        case HTP_PROF_TRACE:
            d->usecs  = HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - d->usecs);
            d->cycles_stop = hex_get_cycles();
            break;
        default:
            break;
    }
}

static int op_fence(struct htp_ops_context * octx) {
    struct htp_context *ctx = octx->ctx;
    struct htp_thread_trace * tr = &ctx->trace[0];
    const uint32_t seq = (uint32_t) octx->op_params[0];

    htp_trace_event_start(tr, HTP_TRACE_EVT_FENCE, (uint16_t) seq);

    const struct htp_tensor * sync = octx->src[0];
    atomic_uint * sync_fence = (atomic_uint *) sync->data;
    uint64_t spins = 0;
    while (1) {
        Q6_dccleaninva_A((void *) sync_fence);
        asm volatile ("syncht" : : : "memory");
        uint32_t val = atomic_load(&sync_fence[0]);
        if ((int32_t)(val - seq) >= 0) {
            break;
        }
        if (++spins > HTP_FENCE_TIMEOUT) {
            FARF(ERROR, "ggml-hex: sync-wait TIMEOUT : fence %p spins %llu seq %u\n", sync_fence, spins, seq);
            break;
        }
        hex_pause();
    }

    htp_trace_event_stop(tr, HTP_TRACE_EVT_FENCE, (uint16_t) seq);

    FARF(HIGH, "ggml-hex: sync-done : fence %p spins %llu seq %u\n", sync_fence, spins, seq);
    return HTP_STATUS_OK;
}

static int execute_op(struct htp_ops_context * octx) {
    switch (octx->op) {
        case HTP_OP_FENCE:
            return op_fence(octx);

        case HTP_OP_ALLREDUCE:
        case HTP_OP_ALLREDUCE_ADD:
            return op_allreduce(octx);

        case HTP_OP_MUL_MAT:
        case HTP_OP_MUL_MAT_ADD:
            return op_matmul(octx);

        case HTP_OP_MUL_MAT_ID:
            return op_matmul_id(octx);

        case HTP_OP_MUL_MAT_NX:
            return op_matmul_nx(octx);

        case HTP_OP_MUL:
        case HTP_OP_ADD:
        case HTP_OP_SUB:
        case HTP_OP_DIV:
        case HTP_OP_ADD_ID:
            return op_binary(octx);

        case HTP_OP_NORM:
        case HTP_OP_RMS_NORM:
        case HTP_OP_RMS_NORM_MUL:
        case HTP_OP_SCALE:
        case HTP_OP_CLAMP:
        case HTP_OP_SQR:
        case HTP_OP_SQRT:
        case HTP_OP_UNARY_SOFTPLUS:
        case HTP_OP_UNARY_SIGMOID:
        case HTP_OP_UNARY_SILU:
        case HTP_OP_UNARY_GELU:
        case HTP_OP_UNARY_NEG:
        case HTP_OP_UNARY_EXP:
        case HTP_OP_UNARY_TANH:
        case HTP_OP_UNARY_ABS:
        case HTP_OP_UNARY_LOG:
        case HTP_OP_L2_NORM:
            return op_unary(octx);

        case HTP_OP_GLU_SWIGLU:
        case HTP_OP_GLU_SWIGLU_OAI:
        case HTP_OP_GLU_GEGLU:
            return op_activations(octx);

        case HTP_OP_SOFTMAX:
            return op_softmax(octx);

        case HTP_OP_ROPE:
            return op_rope(octx);

        case HTP_OP_FLASH_ATTN_EXT:
            return op_flash_attn_ext(octx);

        case HTP_OP_SET_ROWS:
            return op_set_rows(octx);

        case HTP_OP_GET_ROWS:
            return op_get_rows(octx);

        case HTP_OP_SUM_ROWS:
            return op_sum_rows(octx);

        case HTP_OP_CPY:
            return op_cpy(octx);

        case HTP_OP_REPEAT:
            return op_repeat(octx);

        case HTP_OP_ARGSORT:
            return op_argsort(octx);

        case HTP_OP_SSM_CONV:
            return op_ssm_conv(octx);

        case HTP_OP_CUMSUM:
            return op_cumsum(octx);

        case HTP_OP_FILL:
            return op_fill(octx);

        case HTP_OP_DIAG:
            return op_diag(octx);

        case HTP_OP_SOLVE_TRI:
            return op_solve_tri(octx);

        case HTP_OP_PAD:
            return op_pad(octx);

        case HTP_OP_IM2COL:
            return op_im2col(octx);

        case HTP_OP_CONCAT:
            return op_concat(octx);

        case HTP_OP_GATED_DELTA_NET:
            return op_gated_delta_net(octx);

        case HTP_OP_TRI:
            return op_unary(octx);

        case HTP_OP_INVALID:
            break;
    }

    FARF(ERROR, "Unknown Op %u", octx->op);
    return -1;
}

static inline bool reuse_buf(struct htp_context *ctx, uint32_t *m_reuse, struct htp_buf_desc *b) {
    b->base = NULL;

    for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
        struct htp_mmap *m = ctx->mmap + i;
        if (m->size && m->fd == b->fd) {
            b->base   = m->base;
            *m_reuse |= (1 << i);
            return true;
        }
    }

    return false;
}

static inline void drop_mmap(struct htp_context *ctx, struct htp_mmap *m) {
    if (m->size) {
        FARF(ALWAYS, "unmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size);
        htp_munmap((void *) m->base, m->size);
        m->size = 0;
        m->base = 0;
        m->fd   = -1;
    }
}

static inline void mmap_buf(struct htp_context *ctx, struct htp_buf_desc *b) {
    if (b->base) return; // already mapped

    // find unused mapping
    for (uint32_t i=0; i < HTP_MAX_MMAPS; i++) {
        struct htp_mmap *m = &ctx->mmap[i];
        if (!m->size) {
            void *va = htp_mmap(b->fd, b->size);
            if (va == NULL) {
                FARF(ERROR, "mmap failed : fd %u size %u", b->fd, (uint32_t) b->size);
                abort(); // can't do much else at this point
            }

            m->base   = b->base = (uint64_t) va;
            m->fd     = b->fd;
            m->size   = b->size;

            FARF(ALWAYS, "mmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size);
            return;
        }
    }

    FARF(ERROR, "mmap failed : exceeded mapping capacity limit of %u", HTP_MAX_MMAPS);
    abort();
}

static void prep_op_bufs(struct htp_context *ctx, struct htp_buf_desc *bufs, uint32_t n_bufs) {
    uint32_t m_reuse = 0; // mmap reuse mask (index from ctx->mmap array)
    uint32_t b_reuse = 0; // buf reuse count

    uint64_t m_vmem  = 0; // mapped vmem
    uint64_t e_vmem  = 0; // extra  vmem

    // See what we can reuse
    for (uint32_t i=0; i < n_bufs; i++) {
        struct htp_buf_desc *b = bufs + i;
        if (reuse_buf(ctx, &m_reuse, b)) { b_reuse++; } else { e_vmem += b->size; }
        FARF(HIGH, "prep-buf #%u : pass0 fd %u base %p size %u flags 0x%x", i, b->fd, (void*) b->base, (uint32_t) b->size, b->flags);
    }

    if (b_reuse == n_bufs) return; // all bufs reuse existing mappings

    // See how much vmem we have mmaped right now
    for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) { m_vmem += ctx->mmap[i].size; }

    FARF(HIGH, "prep-bufs : pass1 mmap-vmem %zu extra-vmem %zu max-vmem %zu : n-bufs %u b-reuse %u",
            (size_t) m_vmem, (size_t) e_vmem, (size_t) ctx->max_vmem, n_bufs, b_reuse);

    if ((m_vmem + e_vmem) > ctx->max_vmem) {
        // Drop unused mappings
        for (uint32_t i=0; i < HTP_MAX_MMAPS; i++) {
            bool used = m_reuse & (1<<i);
            if (!used) { drop_mmap(ctx, ctx->mmap + i); }
        }
    }

    // Create missing mappings
    for (uint32_t i=0; i < n_bufs; i++) {
        struct htp_buf_desc *b = bufs + i;
        mmap_buf(ctx, b);
        FARF(HIGH, "prep-buf #%u : pass1 fd %u base %p size %u flags 0x%x", i, b->fd, (void*) b->base, (uint32_t) b->size, b->flags);
    }
}

static void prep_tensor(struct htp_context *ctx, struct htp_buf_desc *bufs, struct htp_tensor *tens, uint32_t idx, struct htp_tensor *t) {
    uint32_t offset = t->data;
    uint32_t size   = t->size;
    uint32_t bi     = t->bi;

    t->data  = (uint32_t) (bufs[bi].base + offset);  // update data to the actual pointer

    FARF(HIGH, "prep-tensor #%u: bi %u offset %u size %u data %p : %u:%u:%u:%u", idx, t->bi, offset, t->size, (void*) t->data,
        t->ne[0], t->ne[1], t->ne[3], t->ne[3]);
}

static void prep_tensors(struct htp_context *ctx, struct htp_buf_desc *bufs, struct htp_tensor *tens, uint32_t n_tens) {
    for (uint32_t i=0; i < n_tens; i++) {
        prep_tensor(ctx, bufs, tens, i, tens + i);
    }
}

static int proc_op_req(struct htp_ops_context * octx, struct htp_tensor *tens, uint32_t idx, struct htp_op_desc * op) {
    memcpy(octx->op_params, op->params, sizeof(octx->op_params));
    memcpy(octx->kernel_params, op->kernel_params, sizeof(octx->kernel_params));
    octx->flags = op->flags;
    octx->op    = op->opcode;

    FARF(HIGH, "proc-op #%u: opcode %u flags 0x%x", idx, octx->op, octx->flags);

    // Prep input tensors
    for (uint32_t i=0; i<HTP_OP_MAX_INPUTS; i++) {
        uint16_t src_idx = op->src[i];
        if (src_idx == 0xffff) {
            octx->src[i]     = NULL;
            octx->src_dma[i] = NULL;
            continue;
        }

        struct htp_tensor *src = tens + src_idx;
        octx->src[i]     = src;
        octx->src_dma[i] = octx->ctx->dma; // FIXME: ? octx->ctx->dma_cached : octx->ctx->dma;

        FARF(HIGH, "prep-src #%u: data %p size %u : %u:%u:%u:%u", op->src[i], (void*) src->data, src->size,
            src->ne[0], src->ne[1], src->ne[3], src->ne[3]);
    }

    htp_tensor_flush_all(octx->ctx, octx->src, HTP_OP_MAX_INPUTS);

    // Prep output tensors
    for (uint32_t i = 0; i < HTP_OP_MAX_OUTPUTS; i++) {
        uint16_t dst_idx = op->dst[i];
        if (dst_idx == 0xffff) {
            octx->dsts[i]    = NULL;
            octx->dst_dma[i] = NULL;
            continue;
        }
        struct htp_tensor *dst = tens + dst_idx;
        octx->dsts[i]    = dst;
        octx->dst_dma[i] = octx->ctx->dma; // FIXME: ? octx->ctx->dma_cached : octx->ctx->dma;

        FARF(HIGH, "prep-dst[%u] #%u: data %p size %u : %u:%u:%u:%u", i, dst_idx, (void*) dst->data, dst->size,
            dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
    }

    int status = execute_op(octx);

    htp_tensor_dirty_all(octx->ctx, octx->dsts, HTP_OP_MAX_OUTPUTS);

    octx->src0_spad.src = NULL;
    octx->src1_spad.src = NULL;
    octx->src2_spad.src = NULL;
    octx->src3_spad.src = NULL;
    octx->dst_spad.src  = NULL;

    return status;
}

static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_req * req, const struct dspqueue_buffer * dbuf) {
    dspqueue_t queue = ctx->dsp_queue;
    int err;

    const uint32_t n_bufs = req->n_bufs;
    const uint32_t n_tens = req->n_tensors;
    const uint32_t n_ops  = req->n_ops;

    const uint32_t b_size = sizeof(struct htp_buf_desc)  * n_bufs;
    const uint32_t t_size = sizeof(struct htp_tensor)    * n_tens;
    const uint32_t o_size = sizeof(struct htp_op_desc)   * n_ops;
    const uint32_t p_size = sizeof(struct htp_prof_desc) * n_ops;
    const uint32_t tr_size = (HTP_MAX_NTHREADS + 1) * req->n_traces * sizeof(struct htp_trace_desc);

    if (dbuf->size < b_size + t_size + o_size + p_size + tr_size) {
        FARF(ERROR, "invalid opbatch memory block size %u (req %u)", dbuf->size, b_size + t_size + o_size + p_size + tr_size);
        return;
    }

    FARF(HIGH, "processing opbatch #%u: n-bufs %u n-tensors %u n-ops %u n-traces %u : m-size %u b-size %u t-size %u o-size %u", req->id,
            n_bufs, n_tens, n_ops, req->n_traces, dbuf->size, b_size, t_size, o_size);

    // Setup descriptor pointers
    uint8_t * m_ptr = dbuf->ptr;
    struct htp_buf_desc* bufs = (struct htp_buf_desc*)  m_ptr; m_ptr += b_size;
    struct htp_tensor*   tens = (struct htp_tensor*)    m_ptr; m_ptr += t_size;
    struct htp_op_desc*   ops = (struct htp_op_desc*)   m_ptr; m_ptr += o_size;
    struct htp_prof_desc* pds = (struct htp_prof_desc*) m_ptr;

    struct profile_data batch_prof;
    profile_start(HTP_PROF_BASIC, &batch_prof);

    memset(ctx->trace, 0, sizeof(ctx->trace));
    if (ctx->profiler == HTP_PROF_TRACE) {
        struct htp_trace_desc * trace_events = (struct htp_trace_desc *) (m_ptr + p_size);
        for (int t = 0; t <= HTP_MAX_NTHREADS; t++) {
            ctx->trace[t].events     = &trace_events[t * req->n_traces];
            ctx->trace[t].max_events = req->n_traces;
        }
    }

    // Clean cache at the start of the batch
    htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
    qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
    hex_l2fetch_block(ctx, ctx->footprint);
    memset(ctx->dirty_ranges, 0, sizeof(ctx->dirty_ranges));
    htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);

    htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_BUFF, 0);
    prep_op_bufs(ctx, bufs, n_bufs);
    htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_BUFF, 0);

    prep_tensors(ctx, bufs, tens, n_tens);

    struct htp_ops_context *octx = &ctx->octx;
    memset(octx, 0, sizeof(*octx));
    octx->n_threads = ctx->n_threads;
    octx->ctx       = ctx;

    work_queue_wakeup(ctx->work_queue);
    if (ctx->hmx_queue) {
        hmx_queue_wakeup(ctx->hmx_queue);
    }

    int op_status = HTP_STATUS_OK;
    for (uint32_t i = 0; i < n_ops && op_status == HTP_STATUS_OK; i++) {
        struct profile_data prof;

        profile_start(ctx->profiler, &prof);

        op_status = proc_op_req(octx, tens, i, &ops[i]);

        profile_stop(ctx->profiler, &prof);

        if (ctx->profiler) {
            pds[i].opcode = ops[i].opcode;
            pds[i].usecs  = prof.usecs;
            pds[i].cycles_start = prof.cycles_start;
            pds[i].cycles_stop = prof.cycles_stop;
            for (int j = 0; j < HEX_NUM_PMU_COUNTERS; j++) {
                pds[i].pmu[j] = prof.pmu_counters[j];
            }
        }
    }

    if (ctx->hmx_queue) {
        hmx_queue_suspend(ctx->hmx_queue);
        hmx_queue_flush(ctx->hmx_queue);
    }
    work_queue_suspend(ctx->work_queue);

    // Flush remaining dirty tensors at the end of the batch
    htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
    qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
    htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);

    profile_stop(HTP_PROF_BASIC, &batch_prof);

    struct htp_opbatch_rsp rsp;
    memset(&rsp, 0, sizeof(rsp));
    rsp.id           = req->id;
    rsp.status       = op_status;
    rsp.n_bufs       = n_bufs;
    rsp.n_tensors    = n_tens;
    rsp.n_ops        = n_ops;
    rsp.usecs        = batch_prof.usecs;
    rsp.cycles_start = batch_prof.cycles_start;
    rsp.cycles_stop  = batch_prof.cycles_stop;
    rsp.seq          = req->seq;

    if (ctx->profiler == HTP_PROF_TRACE) {
        for (int t = 0; t <= HTP_MAX_NTHREADS; t++) {
            rsp.n_traces[t] = ctx->trace[t].count;
        }
    }

    struct dspqueue_buffer write_dbuf = *dbuf;
    write_dbuf.flags = DSPQUEUE_BUFFER_FLAG_FLUSH_SENDER | DSPQUEUE_BUFFER_FLAG_INVALIDATE_RECIPIENT;

    err = dspqueue_write(queue, 0, 1, &write_dbuf, sizeof(rsp), (const uint8_t *) &rsp, DSPQUEUE_TIMEOUT_NONE);
    if (err != 0) {
        FARF(ERROR, "dspqueue_write failed: 0x%08x", (unsigned) err);
    }
}

#define DSPQUEUE_READ_TIMEOUT_USEC 5000
#define DSPQUEUE_POLL_TIMEOUT_USEC 100
#define DSPQUEUE_POLL_COUNT        100

static void process_ops(struct htp_context * ctx) {
    dspqueue_t queue = ctx->dsp_queue;
    int err;

    uint32_t poll_count = DSPQUEUE_POLL_COUNT;

    vtcm_acquire(ctx);

    while (!ctx->vtcm_needs_release && !atomic_load(&ctx->killed)) {
        struct htp_opbatch_req req;
        uint32_t r_size = sizeof(req);

        struct dspqueue_buffer dbuf;
        uint32_t n_dbufs = 1;
        uint32_t flags   = 0;

        err = dspqueue_read_noblock(queue, &flags, n_dbufs, &n_dbufs, &dbuf, r_size, &r_size, (uint8_t *) &req);
        if (err == AEE_EWOULDBLOCK) {
            if (--poll_count) {
                qurt_sleep(DSPQUEUE_POLL_TIMEOUT_USEC);
                continue;
            }
            break;
        }

        if (err != 0) {
            FARF(ERROR, "dspqueue_read_noblock failed: 0x%08x", (unsigned) err);
            break;
        }

        if (r_size < sizeof(req) || n_dbufs != 1) {
            FARF(ERROR, "invalid request : size %u n-dbufs %u", r_size, n_dbufs);
            continue;
        }

        // Reset poll count for valid requests
        poll_count = DSPQUEUE_POLL_COUNT;

        process_opbatch(ctx, &req, &dbuf);
    }

    vtcm_release(ctx);
}

static void htp_packet_callback(dspqueue_t queue, int error, void * context) {
    (void) queue;
    (void) error;
    struct htp_handle * h = (struct htp_handle *) context;
    if (h && h->ctx) {
        process_ops(h->ctx);
    }
}

static void htp_main_thread(void * context) {
    struct htp_context * ctx = (struct htp_context *) context;

    FARF(HIGH, "htp-main-thread: started");

    while (!atomic_load(&ctx->killed)) {
        uint32_t flags = 0;
        uint32_t num_buffers = 0;
        uint32_t message_length = 0;

        int err = dspqueue_peek(ctx->dsp_queue, &flags, &num_buffers, &message_length, 50000);
        if (err == 0) {
            process_ops(ctx);
        } else if (err == AEE_EWOULDBLOCK || err == AEE_EEXPIRED) {
            continue;
        } else {
            FARF(ERROR, "dspqueue_peek failed: 0x%08x", (unsigned) err);
            break;
        }
    }

    FARF(HIGH, "htp-main-thread: stopped");
}