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

struct ggml_context;
struct ggml_tensor;

namespace pk {

// A single FastConformer encoder layer (NeMo ConformerLayer), built per-layer
// from the GGUF weights. Mirrors NeMo's ConformerLayer.forward exactly:
//
//   r = x
//   r = r + 0.5 * feed_forward1(norm_feed_forward1(r))   # FFN1 (half-step)
//   r = r + self_attn(norm_self_att(r), pos_emb, mask)    # MHSA (RelPosAttention)
//   r = r + conv(norm_conv(r))                            # Conv module
//   r = r + 0.5 * feed_forward2(norm_feed_forward2(r))    # FFN2 (half-step)
//   out = norm_out(r)
//
// feed_forwardN  = linear2(silu(linear1(x))).
// conv (ConformerConvolution, operates on [d_model, T]):
//   pointwise_conv1 (d->2d, k=1) -> GLU(dim=channel) -> [zero padded time pos]
//   -> depthwise_conv (d->d, k=conv_kernel, groups=d)
//   -> norm -> SiLU -> pointwise_conv2 (d->d, k=1).
// Two conv-module variants, selected by config (byte-identical gating):
//   * conv_norm_type=batch_norm (offline models): inference affine fold from
//     running stats (eps 1e-5); symmetric depthwise pad (k-1)/2 each side.
//   * conv_norm_type=layer_norm (streaming models, e.g. the EOU model): LayerNorm
//     over the channel dim per time-frame (eps 1e-5), gamma/beta read from
//     conv.batch_norm.{weight,bias} (NeMo always names the attr `batch_norm`).
//     conv_causal=true -> causal depthwise pad (left k-1, right 0).
// All norm_* (outside conv) are LayerNorm (eps 1e-5). MHSA reuses pk::RelPosAttention.
//
// 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
// center-pad). It is threaded into RelPosAttention (key/query masking) and used
// to zero padded time positions before the depthwise conv, matching NeMo's
// pad_mask handling. Pass valid_len == T to disable masking.
class ConformerLayer {
public:
    ConformerLayer(const ModelLoader& ml, int layer_idx);

    // GRAPH-BUILDER: append the WHOLE conformer layer (FFN1 + MHSA + conv + FFN2
    // + norm_out) to a SHARED graph `ctx`. `xt` is the layer input tensor [D, T]
    // and `pe` is the positional-encoding tensor [D, pos_len], both ALREADY in
    // the graph. Returns the layer output [D, T] (ggml ne0=D fastest, row-major
    // [T, d_model]). Host-built masks / batch-norm fold constants are fed via
    // pk::graph_input_tensor and registered into `pool` (must outlive compute).
    // This is the unit reused by the fused encoder AND the unit test; computing
    // the entire layer as ONE sub-graph (vs the old 5 sub-graphs) is what lets
    // the fused encoder be a single graph.
    // When att_left/att_right >= 0, the self-attention uses NeMo
    // rel_pos_local_attn (banded, O(T*window)): `pe` must then be the LOCAL
    // positional encoding [d_model, att_left+att_right+1]. Defaults (-1, -1) keep
    // full attention with `pe` = [d_model, 2T-1].
    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;

    // 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].
    // att_left/att_right >= 0 routes self-attention to banded local attention
    // (pe = LOCAL [d_model, att_left+att_right+1]); defaults (-1,-1) = full.
    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,
                                     int att_left = -1, int att_right = -1) const;

    // x: [T, d_model]; pos_emb: [pos_len=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;

    // Same as forward(), but also returns the ConformerConvolution sub-module
    // output (NeMo `layers[i].conv` output, row-major [T, d_model]) for parity
    // localization against the baseline `l0_conv_out`.
    void forward_with_conv(const std::vector<float>& x, int T,
                           const std::vector<float>& pos_emb, int pos_len,
                           int valid_len,
                           std::vector<float>& out,
                           std::vector<float>& conv_out) const;

    // Run JUST the ConformerConvolution sub-module (everything AFTER norm_conv:
    // pointwise_conv1 -> GLU -> [pad_mask] -> depthwise_conv -> norm -> SiLU ->
    // pointwise_conv2) on `conv_in` (= norm_conv(residual), row-major [T,d_model]).
    // Test entry point for validating the conv module — including the layer_norm
    // and causal-conv variants — in ISOLATION from the surrounding attention.
    // `out` is row-major [T, d_model].
    void conv_module_forward(const std::vector<float>& conv_in, int T,
                             int valid_len, std::vector<float>& out) const;

private:
    const ModelLoader& ml_;
    int layer_idx_;
    int d_model_;
    int n_heads_;
    int ff_dim_;
    int conv_kernel_;
    std::string conv_norm_type_;  // "batch_norm" (offline) or "layer_norm" (streaming)
    bool conv_causal_ = false;    // causal depthwise conv pad (left k-1, right 0)
};

} // namespace pk