#pragma clang diagnostic ignored "-Wunused-variable"
#pragma clang diagnostic ignored "-Wunused-function"
#pragma clang diagnostic ignored "-Wunused-but-set-variable"

#include <HAP_farf.h>
#include <HAP_perf.h>
#include <hexagon_protos.h>
#include <hexagon_types.h>
#include <string.h>

#define GGML_COMMON_DECL_C
#include "ggml-common.h"
#include "htp-ctx.h"
#include "htp-ops.h"
#include "hvx-utils.h"
#include "hex-dma.h"
#include "hex-profile.h"
#include "htp-vtcm.h"

struct htp_im2col_context {
    struct htp_ops_context * octx;
    uint32_t                 npatches_per_thread;  // patches = N*OH*OW (pure-DDR kernel)

    uint32_t pe_rows_per_thread;                   // N*OH rows per worker
    uint32_t pe_src_row_bytes;                     // one output row's source: IC*KH*IW*4, rounded 256
    uint32_t pe_dst_row_bytes;                     // one output row's dst: OW*patch_stride*2, rounded 256

    // Patch-embed DMA path VTCM ping-pong.
    uint8_t * pe_vtcm_src;                         // base of the 2x src buffers region
    uint8_t * pe_vtcm_dst;                         // base of the 2x dst buffers region
    uint32_t  pe_src_size_per_thread;              // 2 * pe_src_row_bytes
    uint32_t  pe_dst_size_per_thread;              // 2 * pe_dst_row_bytes
};

// Per-op VTCM layout for the patch-embed DMA path
struct htp_im2col_vtcm_layout {
    size_t off_src;
    size_t off_dst;
    size_t src_bytes_per_thread;
    size_t dst_bytes_per_thread;
    size_t total_bytes;
};

static inline void htp_im2col_vtcm_layout_build(struct htp_im2col_vtcm_layout * L,
                                                size_t                          src_row_bytes,
                                                size_t                          dst_row_bytes,
                                                uint32_t                        n_threads) {
    L->src_bytes_per_thread = 2 * src_row_bytes;
    L->dst_bytes_per_thread = 2 * dst_row_bytes;

    L->off_src     = 0;
    L->off_dst     = L->off_src + L->src_bytes_per_thread * n_threads;
    L->total_bytes = L->off_dst + L->dst_bytes_per_thread * n_threads;
}

#define IM2COL_PATCHEMBED_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG)                        \
    static void FNAME(unsigned int nth, unsigned int ith, void * data) {                                  \
        struct htp_im2col_context * ictx        = (struct htp_im2col_context *) data;                     \
        struct htp_ops_context *    octx        = ictx->octx;                                             \
        struct htp_thread_trace * restrict tr   = &octx->ctx->trace[ith];                                 \
        const struct htp_tensor * restrict src1 = octx->src[1];                                           \
        const struct htp_tensor * restrict dst  = octx->dst;                                              \
        const int32_t  s0                       = octx->op_params[0];                                     \
        const int32_t  s1                       = octx->op_params[1];                                     \
        const int32_t  p0                       = octx->op_params[2];                                     \
        const int32_t  p1                       = octx->op_params[3];                                     \
        const int32_t  d0                       = octx->op_params[4];                                     \
        const int32_t  d1                       = octx->op_params[5];                                     \
        const uint32_t N                        = src1->ne[3];                                            \
        const uint32_t IC                       = src1->ne[2];                                            \
        const uint32_t IH                       = src1->ne[1];                                            \
        const uint32_t IW                       = src1->ne[0];                                            \
        const uint32_t KH                       = octx->src[0]->ne[1];                                    \
        const uint32_t KW                       = octx->src[0]->ne[0];                                    \
        const uint32_t OH                       = dst->ne[2];                                             \
        const uint32_t OW                       = dst->ne[1];                                             \
        const uint32_t patch_stride             = IC * KH * KW;                                           \
        const float * restrict src_data         = (const float *) src1->data;                             \
        DST_CTYPE * restrict dst_data           = (DST_CTYPE *) dst->data;                                \
        const uint32_t npatches                 = N * OH * OW;                                            \
        const uint32_t patch_start              = ictx->npatches_per_thread * ith;                        \
        const uint32_t patch_end                = MIN(patch_start + ictx->npatches_per_thread, npatches); \
        if (patch_start >= patch_end) {                                                                   \
            return;                                                                                       \
        }                                                                                                 \
        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, patch_start);                                   \
        for (uint32_t p = patch_start; p < patch_end; p++) {                                              \
            const uint32_t iow             = p % OW;                                                      \
            const uint32_t ioh             = (p / OW) % OH;                                               \
            const uint32_t in              = p / (OW * OH);                                               \
            DST_CTYPE * restrict dst_patch = dst_data + (uint64_t) p * patch_stride;                      \
            for (uint32_t iic = 0; iic < IC; iic++) {                                                     \
                const float * restrict src_plane = src_data + ((uint64_t) in * IC + iic) * IH * IW;       \
                for (uint32_t ikh = 0; ikh < KH; ikh++) {                                                 \
                    const int32_t iih            = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1;          \
                    DST_CTYPE * restrict out_run = dst_patch + iic * (KH * KW) + ikh * KW;                \
                    if (iih < 0 || iih >= (int32_t) IH) {                                                 \
                        SPLAT_FN(out_run, 0.0f, KW);                                                      \
                        continue;                                                                         \
                    }                                                                                     \
                    const int32_t iiw0             = (int32_t) iow * s0 - p0;                             \
                    const float * restrict src_run = src_plane + (uint64_t) iih * IW + iiw0;              \
                    if (d0 == 1) {                                                                        \
                        /* contiguous source run: [lo,hi) is in-bounds, tails are zero pad */             \
                        const int32_t lo = iiw0 < 0 ? -iiw0 : 0;                                          \
                        int32_t       hi = (int32_t) IW - iiw0;                                           \
                        if (hi > (int32_t) KW) {                                                          \
                            hi = (int32_t) KW;                                                            \
                        }                                                                                 \
                        if (hi <= lo) {                                                                   \
                            SPLAT_FN(out_run, 0.0f, KW);                                                  \
                        } else {                                                                          \
                            if (lo > 0) {                                                                 \
                                SPLAT_FN(out_run, 0.0f, (uint32_t) lo);                                   \
                            }                                                                             \
                            COPY_FN((uint8_t *) (out_run + lo), (const uint8_t *) (src_run + lo),         \
                                    (uint32_t) (hi - lo));                                                \
                            if (hi < (int32_t) KW) {                                                      \
                                SPLAT_FN(out_run + hi, 0.0f, (KW - (uint32_t) hi));                       \
                            }                                                                             \
                        }                                                                                 \
                        continue;                                                                         \
                    }                                                                                     \
                    for (uint32_t ikw = 0; ikw < KW; ikw++) {                                             \
                        const int32_t iiw = (int32_t) iow * s0 + (int32_t) ikw * d0 - p0;                 \
                        out_run[ikw]      = (iiw < 0 || iiw >= (int32_t) IW) ?                            \
                                                (DST_CTYPE) 0.0f :                                        \
                                                (DST_CTYPE) src_plane[(uint64_t) iih * IW + iiw];         \
                    }                                                                                     \
                }                                                                                         \
            }                                                                                             \
        }                                                                                                 \
        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, patch_start);                                    \
    }

IM2COL_PATCHEMBED_BODY(im2col_patchembed_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "f32-f16")
IM2COL_PATCHEMBED_BODY(im2col_patchembed_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "f32-f32")

#define IM2COL_PATCHEMBED_DMA_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG)                               \
    static void FNAME(unsigned int nth, unsigned int ith, void * data) {                                             \
        struct htp_im2col_context * ictx        = (struct htp_im2col_context *) data;                                \
        struct htp_ops_context *    octx        = ictx->octx;                                                        \
        struct htp_thread_trace * restrict tr   = &octx->ctx->trace[ith];                                            \
        const struct htp_tensor * restrict src1 = octx->src[1];                                                      \
        const struct htp_tensor * restrict dst  = octx->dst;                                                         \
        const uint32_t N = src1->ne[3], IC = src1->ne[2], IH = src1->ne[1], IW = src1->ne[0];                        \
        const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0];                                           \
        const uint32_t OH = dst->ne[2], OW = dst->ne[1];                                                             \
        const uint32_t patch_stride     = IC * KH * KW;                                                              \
        const float * restrict src_data = (const float *) src1->data;                                                \
        DST_CTYPE * restrict dst_data   = (DST_CTYPE *) dst->data;                                                   \
        dma_queue *    dmaq             = octx->ctx->dma[ith];                                                       \
        uint8_t *      src_base         = ictx->pe_vtcm_src + ith * ictx->pe_src_size_per_thread;                    \
        uint8_t *      dst_base         = ictx->pe_vtcm_dst + ith * ictx->pe_dst_size_per_thread;                    \
        float *        srcb             = (float *) src_base;                                                        \
        DST_CTYPE *    dstb             = (DST_CTYPE *) dst_base;                                                    \
        const uint32_t nrows            = N * OH;                                                                    \
        const uint32_t per_thread       = ictx->pe_rows_per_thread;                                                  \
        const uint32_t row_start        = per_thread * ith;                                                          \
        const uint32_t row_end          = MIN(row_start + per_thread, nrows);                                        \
        if (row_start >= row_end)                                                                                    \
            return;                                                                                                  \
        for (uint32_t r = row_start; r < row_end; r++) {                                                             \
            const uint32_t in  = r / OH;                                                                             \
            const uint32_t ioh = r % OH;                                                                             \
            for (uint32_t ikh = 0; ikh < KH; ikh++) {                                                                \
                int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh;                                          \
                int     ok  = (iih >= 0 && iih < (int32_t) IH);                                                      \
                for (uint32_t iic = 0; iic < IC; iic++) {                                                            \
                    float *       vdst = srcb + ((uint64_t) (iic * KH + ikh)) * IW;                                  \
                    const float * _vsrc =                                                                            \
                        ok ? (src_data + ((uint64_t) (in * IC + iic) * IH + iih) * IW) : (const float *) vdst;       \
                    dma_queue_push_ddr_to_vtcm(                                                                      \
                        dmaq, dma_make_ptr((uint8_t *) vdst, ok ? (const uint8_t *) _vsrc : (const uint8_t *) vdst), \
                        IW * sizeof(float), IW * sizeof(float), ok ? 1 : 0);                                         \
                }                                                                                                    \
            }                                                                                                        \
            for (uint32_t i = 0; i < IC * KH; i++)                                                                   \
                dma_queue_pop(dmaq);                                                                                 \
            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r);                                                    \
            for (uint32_t iow = 0; iow < OW; iow++) {                                                                \
                DST_CTYPE * dst_patch = dstb + (uint64_t) iow * patch_stride;                                        \
                for (uint32_t ikh = 0; ikh < KH; ikh++) {                                                            \
                    int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh;                                      \
                    for (uint32_t iic = 0; iic < IC; iic++) {                                                        \
                        DST_CTYPE * out_run = dst_patch + iic * (KH * KW) + ikh * KW;                                \
                        if (iih < 0 || iih >= (int32_t) IH) {                                                        \
                            SPLAT_FN(out_run, 0.0f, KW);                                                             \
                            continue;                                                                                \
                        }                                                                                            \
                        const float * src_run = srcb + ((uint64_t) (iic * KH + ikh)) * IW + (uint64_t) iow * KW;     \
                        COPY_FN((uint8_t *) out_run, (const uint8_t *) src_run, KW);                                 \
                    }                                                                                                \
                }                                                                                                    \
            }                                                                                                        \
            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r);                                                     \
            DST_CTYPE * ddr_row = dst_data + ((uint64_t) (in * OH + ioh) * OW) * patch_stride;                       \
            dma_queue_push_vtcm_to_ddr(dmaq, dma_make_ptr((uint8_t *) ddr_row, (uint8_t *) dstb),                    \
                                       OW * patch_stride * (DST_ELEM), OW * patch_stride * (DST_ELEM), 1);           \
            dma_queue_flush(dmaq);                                                                                   \
        }                                                                                                            \
    }

IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_thread,     __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "pe-dma-f16")
IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_f32_thread, float,  hvx_copy_f32_uu,     hvx_splat_f32_u, sizeof(float),  "pe-dma-f32")

static bool im2col_use_patchembed_dma(const struct htp_ops_context * octx) {
    const int32_t s0 = octx->op_params[0], s1 = octx->op_params[1];
    const int32_t p0 = octx->op_params[2], p1 = octx->op_params[3];
    const int32_t d0 = octx->op_params[4], d1 = octx->op_params[5];
    const int     is_2D = octx->op_params[6] == 1;
    if (!is_2D) {
        return false;
    }
    if (octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_F32) {
        return false;
    }
    const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0];
    if (s0 != (int32_t) KW || s1 != (int32_t) KH) {
        return false;  // non-overlapping
    }
    if (p0 != 0 || p1 != 0) {
        return false;  // no padding
    }
    if (d0 != 1 || d1 != 1) {
        return false;  // no dilation
    }
    return true;
}

// Sizes the per-thread 2x(src,dst) VTCM ping-pong for the patch-embed DMA path.
// Returns false if it doesn't fit the VTCM budget (caller falls back).
static bool im2col_patchembed_dma_fits(struct htp_ops_context *    octx,
                                       struct htp_im2col_context * ictx,
                                       uint32_t                    n_threads) {
    const uint32_t IC = octx->src[1]->ne[2], IW = octx->src[1]->ne[0];
    const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0];
    const uint32_t OW           = octx->dst->ne[1];
    const uint32_t patch_stride = IC * KH * KW;

    ictx->pe_src_row_bytes  = hex_round_up(IC * KH * IW * sizeof(float), 256);
    const uint32_t dst_elem = (octx->dst->type == HTP_TYPE_F16) ? sizeof(__fp16) : sizeof(float);
    ictx->pe_dst_row_bytes  = hex_round_up(OW * patch_stride * dst_elem, 256);

    // 2 src + 2 dst buffers per thread (ping-pong), src region first then dst.
    struct htp_im2col_vtcm_layout L;
    htp_im2col_vtcm_layout_build(&L, ictx->pe_src_row_bytes, ictx->pe_dst_row_bytes, n_threads);
    if (L.total_bytes > octx->ctx->vtcm_size) {
        return false;
    }

    uint8_t * const base        = octx->ctx->vtcm_base;
    ictx->pe_vtcm_src           = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src);
    ictx->pe_vtcm_dst           = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
    ictx->pe_src_size_per_thread = (uint32_t) L.src_bytes_per_thread;
    ictx->pe_dst_size_per_thread = (uint32_t) L.dst_bytes_per_thread;
    return true;
}

int op_im2col(struct htp_ops_context * octx) {
    const struct htp_tensor * src1 = octx->src[1];
    const struct htp_tensor * dst  = octx->dst;

    if (src1->type != HTP_TYPE_F32 || (dst->type != HTP_TYPE_F16 && dst->type != HTP_TYPE_F32)) {
        FARF(ERROR, "im2col: only (F32 image -> F16/F32 columns) supported");
        return HTP_STATUS_NO_SUPPORT;
    }

    const uint32_t N         = src1->ne[3];
    const uint32_t OH        = dst->ne[2];
    const uint32_t OW        = dst->ne[1];
    const uint32_t npatches  = N * OH * OW;
    const uint32_t n_threads = MIN(octx->n_threads, npatches);

    if ((octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) || n_threads == 0) {
        return HTP_STATUS_OK;
    }

    struct htp_im2col_context ictx = { 0 };
    ictx.octx                      = octx;
    ictx.npatches_per_thread       = (npatches + n_threads - 1) / n_threads;

    // Clean non-overlapping patch-embed -> DMA kernel (if it fits VTCM);
    // everything else (padding/dilation/stride edges) -> pure-DDR kernel.
    if (im2col_use_patchembed_dma(octx)) {
        const uint32_t nrows = N * OH;
        const uint32_t pth   = MIN(octx->n_threads, nrows);
        if (pth > 0 && im2col_patchembed_dma_fits(octx, &ictx, pth)) {
            ictx.pe_rows_per_thread = (nrows + pth - 1) / pth;
            if (dst->type == HTP_TYPE_F16) {
                work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_thread, &ictx, pth);
            } else {
                work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_f32_thread, &ictx, pth);
            }
            return HTP_STATUS_OK;
        }
        // else: doesn't fit -> fall through to the pure-DDR kernel below.
    }

    if (dst->type == HTP_TYPE_F16) {
        work_queue_run(octx->ctx->work_queue, im2col_patchembed_thread, &ictx, n_threads);
    } else {
        work_queue_run(octx->ctx->work_queue, im2col_patchembed_f32_thread, &ictx, n_threads);
    }
    return HTP_STATUS_OK;
}
