#pragma once
#include "model_loader.hpp"
#include "graph_builder.hpp"
#include <vector>

struct ggml_context;
struct ggml_tensor;

namespace pk {

// Relative-position multi-head self-attention (Transformer-XL style) as used by
// the FastConformer encoder — NeMo RelPositionMultiHeadAttention.
//
// IMPORTANT (module boundary): this mirrors NeMo's `self_attn(...)` exactly. The
// input `x` is the ALREADY-normalized attention input, i.e. NeMo's
// `norm_self_att(residual)` (note: `residual` already includes FFN1, so it is
// NOT `norm_self_att(enc_pre_layers)`). The norm itself lives in the conformer
// layer (next task), NOT here. q, k and v are all projected from this same `x`.
//
// Layout convention (matches the rest of the port and the baseline GGUF):
//   x       row-major [T, d_model]      (d_model fastest)
//   pos_emb row-major [2T-1, d_model]   (d_model fastest)
//   out     row-major [T, d_model]      (d_model fastest)
//
// `valid_len` is the number of non-padding frames. Frames >= valid_len are
// padding: their key columns get -inf in the scores (zero attention weight) and
// their query rows are fully masked, so each padded output row reduces to
// linear_out.bias — exactly matching NeMo's att_mask handling. Pass valid_len ==
// T to disable masking.
class RelPosAttention {
public:
    RelPosAttention(const ModelLoader& ml, int layer_idx);

    // GRAPH-BUILDER: append MHSA ops to a SHARED graph `ctx`. `xt` is the
    // normalized attention input tensor [D, T] (ggml ne0=D fastest) and `pe` is
    // the positional-encoding tensor [D, pos_len], both ALREADY in the graph.
    // Returns the attention output [D, T]. Host-built additive masks are fed via
    // pk::graph_input_tensor and registered into `pool` (must outlive compute).
    // Reused by the fused conformer layer and the unit test.
    //
    // When att_left/att_right >= 0, an additional SYMMETRIC sliding-window mask
    // is applied: query qi may attend to key kj only if -att_left <= qi-kj <=
    // att_right (NeMo rel_pos_local_attn). Defaults (-1, -1) = full context.
    ggml_tensor* build_graph(ggml_context* ctx, ggml_tensor* xt, int T,
                             ggml_tensor* pe, int pos_len, int valid_len,
                             GraphInputPool& pool,
                             int att_left = -1, int att_right = -1) const;

    // LOCAL (banded / Longformer) GRAPH-BUILDER. Appends NeMo rel_pos_local_attn
    // ops to a SHARED graph. `xt` is [D, T] and `pe` is the LOCAL positional
    // encoding [D, att_left+att_right+1]. Each query attends only to keys within
    // [t-att_left, t+att_right] via pad-and-shift, so peak memory is
    // O(T * window) not O(T^2). Returns the attention output [D, T].
    ggml_tensor* build_graph_local(ggml_context* ctx, ggml_tensor* xt, int T,
                                   ggml_tensor* pe, int pos_len, int valid_len,
                                   int att_left, int att_right,
                                   GraphInputPool& pool) const;

    // Batched GRAPH-BUILDER. `xt` is [D, T, B]; `pe` is [D, pos_len] (shared
    // across the batch). `valid_len` is per item (size B). Returns [D, T, B].
    ggml_tensor* build_graph_batched(ggml_context* ctx, ggml_tensor* xt, int T,
                                     int B, ggml_tensor* pe, int pos_len,
                                     const std::vector<int>& valid_len,
                                     GraphInputPool& pool) const;

    // Batched LOCAL (banded / Longformer) GRAPH-BUILDER. Same as
    // build_graph_local but for B>1: `xt` is [D, T, B], `pe` is the LOCAL
    // positional encoding [D, att_left+att_right+1] (shared across the batch),
    // `valid_len` is per item. Returns [D, T, B]. Banded -> O(T*window) memory.
    ggml_tensor* build_graph_batched_local(ggml_context* ctx, ggml_tensor* xt, int T,
                                           int B, ggml_tensor* pe, int pos_len,
                                           const std::vector<int>& valid_len,
                                           int att_left, int att_right,
                                           GraphInputPool& pool) const;

    // CHUNK-MATMUL LOCAL GRAPH-BUILDER. Same math/output as build_graph_local
    // (NeMo rel_pos_local_attn, banded O(T*window)) but built with O(1) graph
    // nodes regardless of window: time is tiled into chunks of `chunk` frames,
    // K/V are gathered as OVER-PADDED overlapping chunks (so ggml's dense
    // ne-product view-bounds check passes), the per-chunk q.k blocks are one
    // batched ggml_mul_mat, and a diagonal "skew" view extracts the [P,T] band.
    // Lets the window go to NeMo's full [128,128] without the pad-and-shift
    // path's O(window) nodes + O(window^2) concat. `pe` is the LOCAL positional
    // encoding [D, att_left+att_right+1]. chunk<=0 picks a default.
    ggml_tensor* build_graph_local_chunked(ggml_context* ctx, ggml_tensor* xt,
                                           int T, ggml_tensor* pe, int pos_len,
                                           int valid_len, int att_left, int att_right,
                                           GraphInputPool& pool, int chunk = -1) const;

    // Thin wrapper over build_graph_local_chunked (test entry point). Same
    // signature/semantics as forward_local; output must match it.
    void forward_local_chunked(const std::vector<float>& x, int T,
                               const std::vector<float>& pos_emb, int pos_len,
                               int valid_len, int att_left, int att_right,
                               std::vector<float>& out, int chunk = -1) const;

    // Batched CHUNK-MATMUL local GRAPH-BUILDER. `xt` is [D, T, B]; `pe` is the
    // LOCAL positional encoding [D, att_left+att_right+1] (shared). Runs the O(1)
    // chunk-matmul construction per item (the 4D chunk kernel can't also carry a
    // batch dim - ggml is 4D), so this is O(B) nodes, still O(1) in the window.
    // `valid_len` is per item. Returns [D, T, B]. Same output as
    // build_graph_batched_local but window-cap-free.
    ggml_tensor* build_graph_batched_local_chunked(ggml_context* ctx, ggml_tensor* xt,
                                                   int T, int B, ggml_tensor* pe, int pos_len,
                                                   const std::vector<int>& valid_len,
                                                   int att_left, int att_right,
                                                   GraphInputPool& pool, int chunk = -1) const;

    // x: [T, d_model]; pos_emb: [2T-1, d_model]; out: [T, d_model].
    void forward(const std::vector<float>& x, int T,
                 const std::vector<float>& pos_emb, int pos_len,
                 int valid_len,
                 std::vector<float>& out) const;

    // Local (banded / Longformer) attention — NeMo rel_pos_local_attn. Each
    // query qi attends only to keys in [qi-att_left, qi+att_right]; pos_emb is
    // the LOCAL positional encoding [att_left+att_right+1, d_model] (i.e. 2W+1
    // for a symmetric [W,W] window), NOT the full 2T-1. Output matches NeMo's
    // Longformer self_attn while bounding memory to O(T * window) instead of
    // O(T^2). x: [T, d_model]; out: [T, d_model].
    void forward_local(const std::vector<float>& x, int T,
                       const std::vector<float>& pos_emb, int pos_len,
                       int valid_len, int att_left, int att_right,
                       std::vector<float>& out) const;
private:
    const ModelLoader& ml_;
    int layer_idx_;
    int d_model_;
    int n_heads_;
    int d_head_;
    // Chunked-limited attention (att_context_style=="chunked_limited"). When set,
    // an extra additive -inf window mask is applied to the scores (see forward()).
    bool chunked_limited_ = false;
    int att_left_  = -1;   // att_context_size[0] (left limit in frames)
    int att_right_ = -1;   // att_context_size[1] (right/lookahead in frames)
};

} // namespace pk