#pragma once

#include "engine/framework/core/execution_context.h"
#include "engine/models/nemotron_asr/assets.h"
#include "engine/models/nemotron_asr/frontend.h"
#include "engine/models/nemotron_asr/weights.h"

#include <cstddef>
#include <cstdint>
#include <memory>
#include <unordered_map>
#include <vector>

namespace engine::models::nemotron_asr {

struct NemotronEncodedAudio {
    std::vector<float> values;
    std::vector<int32_t> durations;
    int64_t frames = 0;
    int64_t valid_frames = 0;
    int64_t hidden_size = 0;
};

struct NemotronEncoderStreamState {
    int64_t attention_seen_frames = 0;
    int64_t attention_cached_frames = 0;
    bool first_chunk = true;
    bool backend_cache_valid = false;
    const void * backend_cache_owner = nullptr;
};

// One speaker's cache-aware stream in NeMo's pad_and_drop_preencoded layout
// (masked speaker-tagged ASR): the K/V and convolution caches live in backend
// tensors owned by this state, so several speakers can share one graph.
struct NemotronEncoderSpeakerCaches;
struct NemotronEncoderSpeakerState {
    std::shared_ptr<NemotronEncoderSpeakerCaches> caches;
    int64_t attention_seen_frames = 0;
    int64_t attention_cached_frames = 0;
};

class NemotronFastConformerEncoderRuntime {
public:
    NemotronFastConformerEncoderRuntime(
        std::shared_ptr<const NemotronASRAssets> assets,
        std::shared_ptr<const NemotronWeights> weights,
        engine::core::ExecutionContext & execution_context,
        size_t graph_arena_bytes);
    ~NemotronFastConformerEncoderRuntime();

    void prepare_capacity(int64_t input_frames, int64_t feature_dim, int64_t lookahead_tokens);
    void prepare_streaming_capacity(int64_t feature_dim, int64_t lookahead_tokens);
    void release_offline_graph();
    NemotronEncodedAudio encode(const NemotronFrontendFeatures & features, int64_t prompt_id, int64_t lookahead_tokens);
    NemotronEncoderStreamState make_stream_state() const;
    NemotronEncodedAudio encode_stream_chunk(
        const NemotronFrontendFeatures & features,
        int64_t prompt_id,
        int64_t lookahead_tokens,
        NemotronEncoderStreamState & state);
    // One pad_and_drop step for several speakers in one graph (batch = rows). Each row
    // holds pre-encode cache frames followed by the new frames; the subsampling runs
    // statelessly over all of them and the first drop_extra_pre_encoded encoder frames
    // are dropped. Rows share the frame count and have their own caches and masks.
    std::vector<NemotronEncodedAudio> encode_pad_and_drop_batch(
        const std::vector<const NemotronFrontendFeatures *> & features,
        int64_t prompt_id,
        int64_t lookahead_tokens,
        const std::vector<NemotronEncoderSpeakerState *> & states);
    int64_t pad_and_drop_cache_frames() const;

private:
    struct Graph;

    void ensure_graph(int64_t input_frames, int64_t feature_dim, int64_t lookahead_tokens);
    Graph & ensure_stream_graph(
        int64_t input_frames,
        int64_t feature_dim,
        int64_t lookahead_tokens,
        int64_t prefix_frames,
        bool first_chunk,
        bool raw_cache = false,
        int64_t rows = 1);
    const std::vector<float> & relative_positional_encoding(int64_t frames);

    std::shared_ptr<const NemotronASRAssets> assets_;
    std::shared_ptr<const NemotronWeights> weights_;
    engine::core::ExecutionContext * execution_context_ = nullptr;
    size_t graph_arena_bytes_ = 0;
    std::unique_ptr<Graph> graph_;
    std::vector<std::unique_ptr<Graph>> stream_graphs_;
    std::vector<float> input_scratch_;
    std::vector<float> output_scratch_;
    std::vector<float> prompt_scratch_;
    std::vector<int32_t> mask_scratch_;
    std::vector<float> attention_mask_scratch_;
    std::unordered_map<int64_t, std::vector<float>> relative_positional_encoding_cache_;
};

}  // namespace engine::models::nemotron_asr