#include "streaming_encoder.hpp"
#include "subsampling.hpp"
#include "pos_enc.hpp"
#include "ggml_graph.hpp"
#include "backend.hpp"
#include "graph_builder.hpp"
#include "ggml.h"
#include <cassert>
#include <cmath>
#include <cstring>
#include <memory>
#include <string>
#include <vector>

namespace pk {

// Weights from the loader are brought into the graph as zero-copy leaves via the
// shared pk::clone_weight / pk::clone_weight_opt (backend.cpp): allowlisted
// linears may be f16/q8_0 and are dequantized by ggml_mul_mat. std::string-name
// overloads keep the existing call sites unchanged.
static ggml_tensor* clone_weight(ggml_context* ctx, const ModelLoader& ml,
                                 const std::string& name) {
    return pk::clone_weight(ctx, ml, name.c_str());
}
static ggml_tensor* clone_weight_opt(ggml_context* ctx, const ModelLoader& ml,
                                     const std::string& name) {
    return pk::clone_weight_opt(ctx, ml, name.c_str());
}

StreamingEncoder::StreamingEncoder(const ModelLoader& ml) : ml_(ml) {
    const ParakeetConfig& c = ml.config();
    d_model_     = (int)c.d_model;
    n_layers_    = (int)c.n_layers;
    n_heads_     = (int)c.n_heads;
    conv_kernel_ = (int)c.conv_kernel;
    xscaling_    = c.xscaling;
    assert(n_heads_ > 0 && d_model_ % n_heads_ == 0);
    d_head_ = d_model_ / n_heads_;
    left_pad_ = conv_kernel_ - 1;   // causal depthwise conv left-context (8)
    assert(c.streaming.present && "StreamingEncoder requires a streaming model");
    const StreamingCfg& s = c.streaming;
    // chunk_size/shift_size/pre_encode_cache_size are stored as the [a,b] arrays;
    // the cache-aware streaming buffer uses the [0] value for the FIRST chunk and
    // the [1] value thereafter (CacheAwareStreamingAudioBuffer.__iter__).
    auto pick = [](const std::vector<int32_t>& v, int idx, int dflt) {
        return (idx < (int)v.size()) ? (int)v[idx] : dflt;
    };
    chunk_first_       = pick(s.chunk_size, 0, 9);
    chunk_main_        = pick(s.chunk_size, 1, 16);
    pre_cache_         = pick(s.pre_encode_cache_size, 1, 9);
    drop_extra_        = s.drop_extra_pre_encoded;
    last_channel_cache_= s.last_channel_cache_size;
    valid_out_len_     = s.valid_out_len;
    att_left_  = c.att_context_left;
    att_right_ = c.att_context_right;
    assert(c.causal_downsampling && "streaming model expects causal subsampling");
    assert(c.conv_causal && "streaming model expects causal depthwise conv");
    assert(c.att_context_style == "chunked_limited");
    reset();
}

void StreamingEncoder::reset() {
    step_ = 0;
    clc_len_ = 0;
    cache_time_.assign(n_layers_,
                       std::vector<float>((size_t)left_pad_ * d_model_, 0.0f));
    cache_channel_.assign(n_layers_,
                          std::vector<float>((size_t)last_channel_cache_ * d_model_, 0.0f));
}

namespace {
// Per-layer streaming graph I/O: the cache reads are graph INPUTS (host data set
// after alloc by graph_input_tensor) and the next-step caches are graph OUTPUTS
// (captured via capture_graph_output and read back after compute). This moves the
// NeMo cache_last_time / cache_last_channel update IN-GRAPH (vs the old per-stage
// host shuffle), so the whole chunk forward is ONE ggml graph.
struct StreamLayerCapture {
    std::vector<float> next_channel;  // updated cache_last_channel [cache_len, D]
    std::vector<float> next_time;     // updated cache_last_time    [LP, D]
};
} // namespace

// Build ONE streaming conformer layer (NeMo ConformerLayer::forward) into the
// SHARED graph `ctx`, threading the per-layer caches in-graph:
//   * MHSA prepends the [cache_len, D] attention K/V cache (update_cache) and uses
//     the streaming chunked-limited + empty-cache mask over [cache_len+Tc] keys.
//   * Conv prepends the [D, LP] depthwise-conv cache before the depthwise conv.
//   * The NEXT caches are emitted as captured outputs (cap->next_channel/_time):
//       next_channel = cat([cache_channel[Tc:], attn_in[:Tc]])  (keep last cache_len)
//       next_time    = last LP cols of [cache_time(LP) ; glu(Tc)]
// `xt` is the chunk input [D, Tc]; `pe` is the pos-emb [D, pos_len] for
// pos_len = 2*(Tc+cache_len)-1. `cache_ch_t` is the K/V cache input [D, cache_len];
// `cache_t_t` is the conv cache input [D, LP]. Returns the layer output [D, Tc].
//
// This MIRRORS the proven (bit-equivalent, max|d| ~7e-6) per-stage layer_step op
// sequence exactly — same ggml ops in the same order — only fused into one graph
// with the cache reads/updates moved in-graph, so the numerics are preserved.
static ggml_tensor* build_stream_layer(
        ggml_context* ctx, const ModelLoader& ml, int layer_idx,
        ggml_tensor* xt, int Tc, ggml_tensor* pe, int pos_len,
        int cache_len, int clc_len, int n_heads, int d_model, int conv_kernel,
        int att_left, int att_right, const std::string& conv_norm_type,
        ggml_tensor* cache_ch_t, ggml_tensor* cache_t_t,
        GraphInputPool& pool, StreamLayerCapture* cap) {
    const int D  = d_model;
    const int H  = n_heads;
    const int dk = D / H;
    const int K  = conv_kernel;
    const int LP = K - 1;
    const float ln_eps = 1e-5f;
    const float scale = 1.0f / std::sqrt((float)dk);
    const int Tk = cache_len + Tc;            // attention key length
    assert(pos_len == 2 * (Tc + cache_len) - 1);

    const std::string pre = "encoder.layers." + std::to_string(layer_idx) + ".";

    auto layer_norm = [&](ggml_tensor* in, const std::string& nm) {
        ggml_tensor* g = clone_weight(ctx, ml, pre + nm + ".weight");
        ggml_tensor* b = clone_weight(ctx, ml, pre + nm + ".bias");
        ggml_tensor* y = ggml_norm(ctx, in, ln_eps);
        y = ggml_mul(ctx, y, g);
        y = ggml_add(ctx, y, b);
        return y;
    };
    auto linear = [&](ggml_tensor* in, const std::string& nm, bool bias) {
        ggml_tensor* W = clone_weight(ctx, ml, pre + nm + ".weight");
        ggml_tensor* y = ggml_mul_mat(ctx, W, in);
        if (bias) { ggml_tensor* B = clone_weight_opt(ctx, ml, pre + nm + ".bias");
                    if (B) y = ggml_add(ctx, y, B); }
        return y;
    };
    auto feed_forward = [&](ggml_tensor* in, const std::string& ff) {
        ggml_tensor* h = linear(in, ff + ".linear1", true);
        h = ggml_silu(ctx, h);
        h = linear(h, ff + ".linear2", true);
        return h;
    };

    // === Stage A: r = x + 0.5*FFN1(norm_ff1(x)); attn_in = norm_self_att(r). ===
    ggml_tensor* h1 = feed_forward(layer_norm(xt, "norm_feed_forward1"),
                                   "feed_forward1");
    h1 = ggml_scale(ctx, h1, 0.5f);
    ggml_tensor* r = ggml_add(ctx, xt, h1);          // [D, Tc]
    ggml_tensor* attn_in = layer_norm(r, "norm_self_att"); // [D, Tc]

    // === Stage B: streaming MHSA (RelPosAttention with attention cache). ===
    // K/V = cat([cache_channel(cache_len), attn_in(Tc)]); Q = attn_in(Tc).
    ggml_tensor* attn_out;
    {
        const std::string ap = pre + "self_attn.";
        ggml_tensor* kvt = ggml_concat(ctx, cache_ch_t, attn_in, /*dim=*/1); // [D, Tk]

        auto lin = [&](const char* w, const char* b, ggml_tensor* in) {
            ggml_tensor* W = clone_weight(ctx, ml, ap + w);
            ggml_tensor* y = ggml_mul_mat(ctx, W, in);
            if (b && ml.tensor(ap + b)) y = ggml_add(ctx, y, clone_weight(ctx, ml, ap + b));
            return y;
        };
        ggml_tensor* q = lin("linear_q.weight", "linear_q.bias", attn_in); // [D, Tc]
        ggml_tensor* k = lin("linear_k.weight", "linear_k.bias", kvt);     // [D, Tk]
        ggml_tensor* v = lin("linear_v.weight", "linear_v.bias", kvt);     // [D, Tk]
        ggml_tensor* p = lin("linear_pos.weight", nullptr, pe);            // [D, P]

        auto to_heads = [&](ggml_tensor* t, int n) {
            t = ggml_reshape_3d(ctx, t, dk, H, n);
            return ggml_cont(ctx, ggml_permute(ctx, t, 0, 2, 1, 3)); // [dk, n, H]
        };
        ggml_tensor* qh = to_heads(q, Tc);
        ggml_tensor* kh = to_heads(k, Tk);
        ggml_tensor* ph = to_heads(p, pos_len);

        ggml_tensor* bu = clone_weight(ctx, ml, ap + "pos_bias_u");
        ggml_tensor* bv = clone_weight(ctx, ml, ap + "pos_bias_v");
        bu = ggml_reshape_3d(ctx, bu, dk, 1, H);
        bv = ggml_reshape_3d(ctx, bv, dk, 1, H);
        ggml_tensor* qu = ggml_add(ctx, qh, bu); // [dk, Tc, H]
        ggml_tensor* qv = ggml_add(ctx, qh, bv); // [dk, Tc, H]

        // ac = q_u @ k^T -> [Tk, Tc, H].
        ggml_tensor* ac = ggml_mul_mat(ctx, kh, qu);

        // bd = q_v @ p^T -> [P, Tc, H]; rel_shift; slice first Tk pos cols.
        ggml_tensor* bd = ggml_mul_mat(ctx, ph, qv); // [P, Tc, H]
        bd = ggml_pad_ext(ctx, bd, 1, 0, 0, 0, 0, 0, 0, 0);      // [P+1, Tc, H]
        bd = ggml_reshape_3d(ctx, bd, Tc, pos_len + 1, H);        // [Tc, P+1, H]
        bd = ggml_view_3d(ctx, bd, Tc, pos_len, H,
                          bd->nb[1], bd->nb[2], bd->nb[1]);       // drop first row
        bd = ggml_cont(ctx, bd);
        bd = ggml_reshape_3d(ctx, bd, pos_len, Tc, H);            // [P, Tc, H]
        bd = ggml_view_3d(ctx, bd, Tk, Tc, H, bd->nb[1], bd->nb[2], 0); // first Tk pos
        bd = ggml_cont(ctx, bd);

        ggml_tensor* scores = ggml_add(ctx, ac, bd); // [Tk, Tc, H]

        // Streaming attention mask [Tk(key), Tc(query)] (0 visible / -inf):
        //   global key index   gk = kj            in [0, Tk)
        //   global query index gq = cache_len + qi in [cache_len, cache_len+Tc)
        //   (1) empty-cache mask: cache cols [0, cache_len-clc_len) are not yet
        //       filled -> masked.
        //   (2) chunked_limited (chunk=att_right+1, left=att_left/chunk):
        //       visible iff 0 <= gq/chunk - gk/chunk <= left_chunks.
        const int chunk = att_right + 1;
        const int left_chunks = (chunk > 0) ? (att_left / chunk) : 0;
        const int empty_cache = cache_len - clc_len;
        std::vector<float>& mask_host = pool.alloc_f32((size_t)Tk * Tc);
        {
            float* md = mask_host.data();
            const float ninf = -INFINITY;
            for (int qi = 0; qi < Tc; ++qi) {
                const int gq = cache_len + qi;
                const int cq = gq / chunk;
                for (int kj = 0; kj < Tk; ++kj) {
                    bool vis = (kj >= empty_cache);
                    if (vis) {
                        const int ck = kj / chunk;
                        const int diff = cq - ck;
                        vis = (diff >= 0 && diff <= left_chunks);
                    }
                    md[(size_t)qi * Tk + kj] = vis ? 0.0f : ninf;
                }
            }
        }
        int64_t mask_ne[2] = {Tk, Tc};
        ggml_tensor* mask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2,
                                mask_ne, mask_host.data(),
                                mask_host.size() * sizeof(float));
        ggml_tensor* attn = ggml_soft_max_ext(ctx, scores, mask, scale, 0.0f);

        ggml_tensor* vh = to_heads(v, Tk);           // [dk, Tk, H]
        ggml_tensor* vtk = ggml_cont(ctx, ggml_permute(ctx, vh, 1, 0, 2, 3)); // [Tk, dk, H]
        ggml_tensor* ctxh = ggml_mul_mat(ctx, vtk, attn); // [dk, Tc, H]
        ggml_tensor* merged = ggml_cont(ctx, ggml_permute(ctx, ctxh, 0, 2, 1, 3));
        merged = ggml_reshape_2d(ctx, merged, D, Tc);
        attn_out = lin("linear_out.weight", "linear_out.bias", merged); // [D, Tc]
    }
    r = ggml_add(ctx, r, attn_out);   // residual after MHSA -> norm_conv -> conv

    // Update attention cache (NeMo update_cache, cache_drop_size=0):
    //   next = cat([cache[Tc:], attn_in[:Tc]]) keep last cache_len.
    // In-graph: drop the oldest Tc cache cols (view starting at col Tc), then
    // concat the Tc current frames; capture for readback into cache_channel_.
    if (cap) {
        const int keep = cache_len - Tc;
        ggml_tensor* next_ch;
        if (keep > 0) {
            ggml_tensor* kept = ggml_view_2d(ctx, cache_ch_t, D, keep,
                                    cache_ch_t->nb[1], (size_t)Tc * cache_ch_t->nb[1]);
            next_ch = ggml_concat(ctx, kept, attn_in, /*dim=*/1); // [D, cache_len]
        } else {
            // Tc >= cache_len: keep only the last cache_len cols of attn_in.
            next_ch = ggml_view_2d(ctx, attn_in, D, cache_len, attn_in->nb[1],
                                   (size_t)(Tc - cache_len) * attn_in->nb[1]);
        }
        next_ch = ggml_cont(ctx, next_ch);
        pk::capture_graph_output(next_ch, &cap->next_channel);
    }

    // === Stage C: streaming conv module (causal depthwise + conv cache). ===
    ggml_tensor* c = layer_norm(r, "norm_conv"); // [D, Tc]
    ggml_tensor* pw1w = clone_weight(ctx, ml, pre + "conv.pointwise_conv1.weight");
    pw1w = ggml_reshape_2d(ctx, pw1w, D, 2 * D);
    ggml_tensor* pw1b = clone_weight_opt(ctx, ml, pre + "conv.pointwise_conv1.bias");
    ggml_tensor* y = ggml_mul_mat(ctx, pw1w, c); // [2D, Tc]
    if (pw1b) y = ggml_add(ctx, y, pw1b);
    ggml_tensor* a = ggml_view_2d(ctx, y, D, Tc, y->nb[1], 0);
    ggml_tensor* b = ggml_view_2d(ctx, y, D, Tc, y->nb[1], (size_t)D * y->nb[0]);
    ggml_tensor* glu = ggml_cont(ctx, ggml_mul(ctx, ggml_cont(ctx, a),
                        ggml_sigmoid(ctx, ggml_cont(ctx, b)))); // [D, Tc]

    // dw_in [D, LP+Tc] = [cache_time(LP) ; glu(Tc)].
    ggml_tensor* dw_in = ggml_concat(ctx, cache_t_t, glu, /*dim=*/1); // [D, LP+Tc]

    // Update conv cache: last LP cols of [cache_time(LP) ; glu(Tc)] (NeMo
    // new_x[:, :, -LP:]); capture for readback into cache_time_.
    if (cap) {
        ggml_tensor* next_t = ggml_view_2d(ctx, dw_in, D, LP, dw_in->nb[1],
                                  (size_t)Tc * dw_in->nb[1]);  // last LP cols
        next_t = ggml_cont(ctx, next_t);
        pk::capture_graph_output(next_t, &cap->next_time);
    }

    ggml_tensor* glu_tc = ggml_cont(ctx, ggml_transpose(ctx, dw_in)); // [LP+Tc, D]
    ggml_tensor* dww = clone_weight(ctx, ml, pre + "conv.depthwise_conv.weight"); // [K,1,C]
    ggml_tensor* dw;
    {
        ggml_tensor* nb = ggml_reshape_4d(ctx, glu_tc, glu_tc->ne[0], 1, glu_tc->ne[1], 1);
        ggml_tensor* ic = ggml_im2col(ctx, dww, nb, 1, 0, 0, 0, 1, 0,
                                      /*is_2D*/false, GGML_TYPE_F32);
        ggml_tensor* r2 = ggml_mul_mat(ctx, ic, dww);
        dw = ggml_reshape_3d(ctx, r2, r2->ne[0], r2->ne[2], 1); // [Tc, C, 1]
    }
    dw = ggml_reshape_2d(ctx, dw, Tc, D);                 // [Tc, C]
    ggml_tensor* dwb = clone_weight_opt(ctx, ml, pre + "conv.depthwise_conv.bias");
    ggml_tensor* dwt = ggml_cont(ctx, ggml_transpose(ctx, dw)); // [C, Tc]
    if (dwb) dwt = ggml_add(ctx, dwt, dwb);

    ggml_tensor* normed;
    if (conv_norm_type == "layer_norm") {
        ggml_tensor* g = clone_weight(ctx, ml, pre + "conv.batch_norm.weight");
        ggml_tensor* bb = clone_weight(ctx, ml, pre + "conv.batch_norm.bias");
        normed = ggml_norm(ctx, dwt, ln_eps);
        normed = ggml_mul(ctx, normed, g);
        normed = ggml_add(ctx, normed, bb);
    } else {
        std::vector<float>& sc_host = pool.alloc_f32(D);
        std::vector<float>& sh_host = pool.alloc_f32(D);
        std::vector<float> g, bb, mm, var;
        pk::weight_to_host_f32(ml, (pre + "conv.batch_norm.weight").c_str(), g);
        pk::weight_to_host_f32(ml, (pre + "conv.batch_norm.bias").c_str(), bb);
        pk::weight_to_host_f32(ml, (pre + "conv.batch_norm.running_mean").c_str(), mm);
        pk::weight_to_host_f32(ml, (pre + "conv.batch_norm.running_var").c_str(), var);
        for (int cc2 = 0; cc2 < D; ++cc2) {
            sc_host[cc2] = g[cc2] / std::sqrt(var[cc2] + 1e-5f);
            sh_host[cc2] = bb[cc2] - mm[cc2] * sc_host[cc2];
        }
        int64_t d_ne[1] = {D};
        ggml_tensor* sc_t = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 1,
                                d_ne, sc_host.data(), sc_host.size() * sizeof(float));
        ggml_tensor* sh_t = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 1,
                                d_ne, sh_host.data(), sh_host.size() * sizeof(float));
        normed = ggml_add(ctx, ggml_mul(ctx, dwt, sc_t), sh_t);
    }
    normed = ggml_silu(ctx, normed);
    ggml_tensor* pw2w = clone_weight(ctx, ml, pre + "conv.pointwise_conv2.weight");
    pw2w = ggml_reshape_2d(ctx, pw2w, D, D);
    ggml_tensor* pw2b = clone_weight_opt(ctx, ml, pre + "conv.pointwise_conv2.bias");
    ggml_tensor* cout = ggml_mul_mat(ctx, pw2w, normed); // [D, Tc]
    if (pw2b) cout = ggml_add(ctx, cout, pw2b);
    r = ggml_add(ctx, r, cout);    // r = r + conv_out

    // === Stage D: r = r + 0.5*FFN2(norm_ff2(r)); out = norm_out(r). ===
    ggml_tensor* h2 = feed_forward(layer_norm(r, "norm_feed_forward2"),
                                   "feed_forward2");
    h2 = ggml_scale(ctx, h2, 0.5f);
    r = ggml_add(ctx, r, h2);
    r = layer_norm(r, "norm_out");
    return r; // [D, Tc]
}

std::vector<float> StreamingEncoder::step(const std::vector<float>& mel_chunk_frames,
                                          int n_mel_frames, bool is_last,
                                          int& n_valid_out) {
    const int D = d_model_;
    const int n_mels = (int)ml_.config().n_mels;
    assert((int)mel_chunk_frames.size() == n_mels * n_mel_frames);

    // ---- 1. Compute Tc / drop host-side (deterministic from the subsampling
    //         arithmetic) so the fused graph knows the shapes up front. The
    //         subsampling builder itself runs in-graph below; here we only
    //         re-derive its output time length (3 stride-2 k=3 causal stages),
    //         matching Subsampling's internal calc. ----
    Subsampling sub(ml_);
    int Tsub_calc = n_mel_frames;
    {
        const int all_paddings = 3; // causal
        int t = n_mel_frames;
        for (int st = 0; st < 3; ++st) t = (t + all_paddings - 3) / 2 + 1;
        Tsub_calc = t;
    }
    const int drop = (step_ != 0) ? drop_extra_ : 0;
    int Tc = Tsub_calc - drop;
    if (Tc < 0) Tc = 0;

    if (Tc == 0) {     // no output this step (e.g. a too-short tail chunk)
        n_valid_out = 0;
        step_ += 1;
        return {};
    }

    // ---- 2. Positional encoding. NeMo pos_enc(cache_len): input_len = Tc +
    //         cache_len; pos_emb spans 2*input_len-1 positions. cache_len = 70. ----
    const int cache_len = last_channel_cache_;
    const int Pn = Tc + cache_len;
    const int pos_len = 2 * Pn - 1;
    const int clc_len_now = clc_len_;   // snapshot for the in-graph mask

    const std::string cnt = ml_.config().conv_norm_type;
    StreamLayerCapture caps[/*max layers*/64];
    assert(n_layers_ <= 64);

    // ---- 3. The WHOLE chunk forward is ONE ggml graph: subsampling (with the
    //         pre-encode-cache window) -> drop -> xscaling -> N streaming conformer
    //         layers (each reading its conv/attn cache as inputs and emitting the
    //         updated caches as captured outputs) -> valid output slice. Computed
    //         in a SINGLE Backend::compute (vs the old ~5 graphs/layer). ----
    std::vector<float> out_full;  // [Tc, D] row-major (the final layer output)
    GraphInputPool pool;
    bool ok = pk::run_graph(/*mem_bytes*/0, /*n_threads*/0,
        [&](ggml_context* ctx) -> ggml_tensor* {
            // Subsampling on the full mel window (already includes pre-encode
            // overlap): the WHOLE window is real audio (no center-pad), so the
            // entry valid length is n_mel_frames. Output [D, Tsub].
            int Tp = 0, valid = 0;
            ggml_tensor* sx = sub.build_graph(ctx, mel_chunk_frames, n_mels,
                                              n_mel_frames, pool, Tp, valid,
                                              /*in_valid_frames*/n_mel_frames);
            assert(Tp == Tsub_calc);

            // Drop drop_extra_pre_encoded leading subsampled frames (NeMo
            // forward_internal; only when cache is present == every step after the
            // first). View columns [drop, Tp) of sx [D, Tp] -> [D, Tc].
            ggml_tensor* x = sx;
            if (drop > 0) {
                x = ggml_view_2d(ctx, sx, D, Tc, sx->nb[1],
                                 (size_t)drop * sx->nb[1]);
                x = ggml_cont(ctx, x);
            }

            // xscaling (gated; off for this model).
            if (xscaling_) {
                x = ggml_scale(ctx, x, std::sqrt((float)D));
            }

            // pos_emb [D, pos_len] fed as a graph input (host-computed once).
            std::vector<float>& pe_host = pool.alloc_f32();
            rel_pos_encoding(Pn, D, pe_host); // row-major [pos_len, D]
            int64_t pe_ne[2] = {D, pos_len};
            ggml_tensor* pe = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, pe_ne,
                                  pe_host.data(), pe_host.size() * sizeof(float));

            // Conformer layer stack (conv + attention caches threaded in-graph).
            const int LP = left_pad_;
            for (int i = 0; i < n_layers_; ++i) {
                // Cache reads as graph inputs (host data set after alloc). The
                // host buffers ARE the live cache vectors, alive across compute.
                int64_t ch_ne[2] = {D, cache_len};
                ggml_tensor* cache_ch_t = pk::graph_input_tensor(ctx,
                        GGML_TYPE_F32, 2, ch_ne, cache_channel_[i].data(),
                        cache_channel_[i].size() * sizeof(float));
                int64_t t_ne[2] = {D, LP};
                ggml_tensor* cache_t_t = pk::graph_input_tensor(ctx,
                        GGML_TYPE_F32, 2, t_ne, cache_time_[i].data(),
                        cache_time_[i].size() * sizeof(float));

                x = build_stream_layer(ctx, ml_, i, x, Tc, pe, pos_len,
                                       cache_len, clc_len_now, n_heads_, D,
                                       conv_kernel_, att_left_, att_right_, cnt,
                                       cache_ch_t, cache_t_t, pool, &caps[i]);
            }
            return x; // [D, Tc] -> row-major [Tc, D]
        }, out_full);
    assert(ok && "streaming fused encoder graph failed"); (void)ok;

    // ---- 4. Commit the captured next-step caches (read back from the graph). ----
    for (int i = 0; i < n_layers_; ++i) {
        assert((int)caps[i].next_channel.size() == cache_len * D);
        assert((int)caps[i].next_time.size() == left_pad_ * D);
        cache_channel_[i].swap(caps[i].next_channel);
        cache_time_[i].swap(caps[i].next_time);
    }
    // After the layer stack the attention cache_last_channel_len has grown by Tc.
    clc_len_ = std::min(cache_len, clc_len_ + Tc);
    step_ += 1;

    // ---- 5. Slice valid output. keep_all_outputs == is_last: mid-stream keep
    //         only valid_out_len frames; final chunk keep all Tc frames. ----
    int valid = is_last ? Tc : std::min(valid_out_len_, Tc);
    n_valid_out = valid;
    std::vector<float> out((size_t)valid * D);
    std::memcpy(out.data(), out_full.data(), (size_t)valid * D * sizeof(float));
    return out;
}

} // namespace pk