#pragma once

#include "ggml-backend-impl.h"
#include "ggml-backend.h"
#include "ggml.h"
#include "openvino/decoder.h"

#include <cstdint>
#include <cstring>
#include <map>
#include <memory>
#include <openvino/core/partial_shape.hpp>
#include <optional>
#include <set>
#include <string>
#include <vector>

struct ModelParams {
    int ctx = -1;
    int ctx_per_seq = -1;
    int ctx_per_seq_swa = -1;
    int n_seq = 1;
    int n_heads_kv = -1;
    int head_size = -1;
    int state_size = -1;  // for SSM molels, eg qwen35
    int32_t rope_params[15];
    bool mixed_rope_params = false;
    std::vector<int> swa_layers;

    std::vector<std::string> kv_names;
    size_t kv_buffer_ctx_id = 0;

    bool same_rope_params(const ModelParams & other) const {
        return mixed_rope_params == other.mixed_rope_params &&
               memcmp(rope_params, other.rope_params, sizeof(int32_t) * 15) == 0;
    }

    bool can_reuse_dynamically(const ModelParams & other) const { return same_rope_params(other); }

    bool can_reuse_statically(const ModelParams & other) const { return same_rope_params(other) && ctx == other.ctx; }

    bool kv_buffer_changed(const ModelParams & other) const { return kv_buffer_ctx_id != other.kv_buffer_ctx_id; }
};

struct ComputeParams {
    int n_seq_active = 1;
    int seq_active_start = 0;
    int attention_size = -1;
    int attention_size_swa = -1;
    int input_len = -1;
    int token_len_per_seq = -1;
    int past_kv_len = -1;
    int output_len = 1;

    int cache_rs_reset_idx = -1;
    int cache_rs_reset_len = -1;
    // SSM/DeltaNet models otionally clear cache_r and cache_s of certain slots in the cgraph
    // 3: [ 18432,     4,     1,     1] RESHAPE              cache_r_l0 (reshaped)
    //    [ 18432,     4,     1,     1]            0: NONE        cache_r_l0
    // 4: [ 18432,     1,     1,     1] VIEW                 cache_r_l0 (reshaped) (view)
    //    [ 18432,     4,     1,     1]            0: RESHAPE     cache_r_l0 (reshaped)
    // 5: [ 18432,     1,     1,     1] SCALE                cache_r_l0 (reshaped) (view) (view)
    //    [ 18432,     1,     1,     1]            0: VIEW        cache_r_l0 (reshaped) (view)

    int s_copy_active_slot_len = -1;
    // SSM/DeltaNet models otionally reorder slots of state cache, to make the active slots contiguous
    // leaf_5 is the inp->s_copy in llama-graph.cpp, eg if there are 8 slots in total and slot 3 and 7
    // are active in the current batch, leaf_5 will be [3, 7, 5, 6, 4]
    //  6: [     2,     1,     1,     1] VIEW                  (view)
    //      [     2,     1,     1,     1]            0: NONE        leaf_5
    //  7: [ 18432,     2,     1,     1] GET_ROWS             conv_states-0
    //      [ 18432,     4,     1,     1]            0: RESHAPE     cache_r_l0 (reshaped)
    //      [     2,     1,     1,     1]            1: VIEW         (view)
    //  8: [     0,     1,     1,     1] VIEW                  (view)
    //      [     2,     1,     1,     1]            0: NONE        leaf_5
    //  9: [ 18432,     0,     1,     1] GET_ROWS             node_9
    //      [ 18432,     4,     1,     1]            0: RESHAPE     cache_r_l0 (reshaped)
    //      [     0,     1,     1,     1]            1: VIEW         (view)
    // 10: [ 18432,     0,     1,     1] VIEW                 cache_r_l0 (view)
    //      [ 18432,     4,     1,     1]            0: NONE        cache_r_l0
    // 11: [ 18432,     0,     1,     1] CPY                  cache_r_l0 (view) (copy of )
    //      [ 18432,     0,     1,     1]            0: GET_ROWS    node_9
    //      [ 18432,     0,     1,     1]            1: VIEW        cache_r_l0 (view)

    struct RsWriteback {
        int slot_begin = 0;  // first cache slot written by the CPY
        int src_begin = 0;   // where the copied data starts in the source tensor (in rows of it)
    };

    std::map<std::string, RsWriteback> rs_writebacks;
    // Offsets of the state cache writeback CPY nodes, keyed by node name. They change with the
    // batch (kv head, active sequence count, token count) and, with rollback enabled
    // (cparams.n_rs_seq > 0), the conv state is written back once per snapshot slot, each snapshot
    // taking a different conv_input window. Passed to the cached model as runtime inputs.
};

class GgmlOvDecoder : public ov::frontend::ggml::GgmlDecoder {
public:
    struct NodeInfo {
        ggml_tensor * node;
        std::string node_name;
        std::string node_op_type;
        std::map<std::string, ggml_tensor *> node_inputs;
        std::map<std::string, std::vector<std::pair<std::string, ggml_tensor *>>> node_inputs_views;
        std::vector<std::string> node_inputs_names;
        int node_op_case = 0;
        void * data_addr;
    };

    // Graph decoder
    GgmlOvDecoder(ggml_cgraph * cgraph,
                  ModelParams & model_params,
                  ComputeParams & compute_params,
                  std::map<std::string, std::shared_ptr<ov::Node>> & model_weights,
                  bool is_static,
                  bool is_stateful = false,
                  bool model_is_splitted = false,
                  bool is_prefill = false,
                  int prefill_chunk_size = 256);

    // Naive graph decoder
    GgmlOvDecoder(ggml_cgraph * cgraph, std::map<std::string, std::shared_ptr<ov::Node>> & model_weights);

    virtual ov::Any get_attribute(const std::string & name) const override {
        return nullptr;
        GGML_UNUSED(name);
    }

    virtual ov::PartialShape get_input_shape(int node_idx, const std::string & name) const override;

    virtual std::vector<size_t> get_input_stride(int node_idx, const std::string & name) const override;

    virtual size_t get_view_input_size(int node_idx, const std::string & name) const override;

    virtual size_t get_view_input_offset(int node_idx, const std::string & name, size_t view_index) const override;

    virtual size_t get_view_input_src_offset(int node_idx, const std::string & name, size_t view_index) const override;

    virtual std::vector<size_t> get_view_input_stride(int node_idx,
                                                      const std::string & name,
                                                      size_t view_index) const override;

    virtual std::vector<size_t> get_view_input_src_stride(int node_idx,
                                                          const std::string & name,
                                                          size_t view_index) const override;

    virtual ov::Shape get_view_input_ggml_shape(int node_idx,
                                                const std::string & name,
                                                size_t view_index) const override;

    virtual ov::Shape get_view_input_src_ggml_shape(int node_idx,
                                                    const std::string & name,
                                                    size_t view_index) const override;

    virtual ov::PartialShape get_view_input_ov_shape(int node_idx,
                                                     const std::string & name,
                                                     size_t view_index) const override;

    virtual ov::PartialShape get_view_input_src_ov_shape(int node_idx,
                                                         const std::string & name,
                                                         size_t view_index) const override;

    virtual std::string get_view_input_name(int node_idx, const std::string & name, size_t view_index) const override;

    virtual std::string get_view_input_src_name(int node_idx,
                                                const std::string & name,
                                                size_t view_index) const override;

    virtual ov::element::Type get_input_type(int node_idx, const std::string & name) const override;

    virtual size_t get_input_size() const override;

    virtual size_t get_input_size(int node_idx) const override;

    virtual void get_input_node(size_t input_port_idx,
                                std::string & producer_name,
                                std::string & producer_output_port_name,
                                size_t & producer_output_port_index) const override {
        GGML_UNUSED(input_port_idx);
        GGML_UNUSED(producer_name);
        GGML_UNUSED(producer_output_port_name);
        GGML_UNUSED(producer_output_port_index);
    }

    virtual std::vector<std::string> get_input_names(int node_idx) const override;

    virtual ov::PartialShape get_output_shape(int node_idx) const override;

    virtual ov::element::Type get_output_type(int node_idx) const override;

    virtual std::vector<size_t> get_output_stride(int node_idx) const override;

    virtual int32_t * get_input_op_params(int node_idx, const std::string & name) const override;

    virtual int32_t * get_output_op_params(int node_idx) const override;

    virtual size_t get_output_op_offset(int node_idx) const override;

    virtual std::vector<std::string> get_output_names(int node_idx) const override;

    virtual std::string get_inplace_op_src(int node_idx) const override;

    virtual bool is_view_like_alias_of(int node_idx, const std::string & view_src_name) const override;

    virtual const std::string & get_op_type() const override;

    virtual const std::string & get_op_type(int node_idx) const override;

    virtual const std::string & get_op_name() const override;

    virtual const std::string & get_op_name(int node_idx) const override;

    virtual int32_t get_op_dynamic_dim(int node_idx) const override;

    virtual void visit_subgraph(
        std::function<void(std::shared_ptr<GgmlDecoder>, int node_idx)> node_visitor) const override;

    ggml_tensor * get_input_ggml_tensor(const std::string & name) const { return m_inputs.at(name); }

    virtual int get_op_case(int node_idx) const override { return m_node_info_list[node_idx].node_op_case; }

    virtual const std::map<std::string, ov::frontend::ggml::ModelInputInfo> & get_model_inputs() const override {
        return m_model_inputs;
    }

    virtual const std::map<std::string, ov::frontend::ggml::ModelExtraInputInfo> & get_model_extra_inputs() const override {
        return m_model_extra_inputs;
    }

    virtual const std::map<std::string, std::shared_ptr<ov::Node>> & get_model_weights() const override {
        return m_model_weights;
    }

    virtual std::set<std::string> get_model_output_names() const override { return m_model_output_names; }

    const std::map<std::string, ggml_tensor *> & get_model_outputs() const { return m_model_outputs; }

    virtual int get_ctx_size() const { return m_model_params.ctx; }

    virtual int get_ctx_per_seq() const { return m_model_params.ctx_per_seq; }

    virtual int get_ctx_per_seq_swa() const { return m_model_params.ctx_per_seq_swa; }

    virtual int get_n_seq() const { return m_model_params.n_seq; }

    virtual int is_swa_layer(int layer) const override {
        return std::find(m_model_params.swa_layers.begin(), m_model_params.swa_layers.end(), layer) !=
               m_model_params.swa_layers.end();
    }

    int get_past_kv_len() const { return m_compute_params.past_kv_len; }

    int get_input_len() const { return m_compute_params.input_len; }

    virtual int32_t * get_rope_params() const override { return const_cast<int32_t *>(m_model_params.rope_params); }

    virtual bool has_mixed_rope_params() const override { return m_model_params.mixed_rope_params; }

    virtual int get_ssm_state_size() const override { return m_model_params.state_size; }

    virtual std::map<std::string, std::string> get_kv_param_res_names() const override;

    virtual bool is_static() const override { return m_is_static; }

    virtual bool is_stateful() const override { return m_is_stateful; }

    int get_static_n_tokens() const { return m_is_prefill ? m_prefill_chunk_size : 1; }

    virtual bool is_splited_model() const override { return m_model_is_splitted; }

    ov::PartialShape get_graph_input_shape(const ggml_tensor * op,
                                           const ggml_tensor * input,
                                           int dynamic_dim_index = -1) const;

    static void dump_cgraph(const ggml_cgraph * cgraph, std::string & filename);

    static std::shared_ptr<ov::Node> create_weight_node(ggml_tensor * tensor, bool naive = false);

    static std::map<std::string, std::shared_ptr<ov::Node>> create_weight_nodes(ggml_cgraph * cgraph,
                                                                                bool naive = false);

    // Collect just the set of weight-tensor names referenced by the graph, without
    // building (or requantizing) any OV weight nodes. Used by topology checks like
    // is_model_splitted that only need name membership.
    static std::set<std::string> collect_weight_names(ggml_cgraph * cgraph);

    const ggml_tensor * get_tensor_used_op(const ggml_tensor * tensor) const;

    const ggml_tensor * get_tensor_from_name(const std::string & name) const;

    void clear_model_weights() { m_model_weights.clear(); }

    static std::pair<ModelParams, ComputeParams> compute_llm_params(ggml_cgraph * cgraph, bool is_static);

    ModelParams get_model_params() const { return m_model_params; }

    ComputeParams get_compute_params() const { return m_compute_params; }

    void set_model_params(const ModelParams & model_params) { m_model_params = model_params; }

    void set_compute_params(const ComputeParams & compute_params) { m_compute_params = compute_params; }

    bool m_is_static = false;
    bool m_is_stateful = false;
    bool m_is_prefill = false;
    bool m_naive = false;
    int m_prefill_chunk_size = 0;
    bool m_model_is_splitted = false;  // label the cgraph is splited or not

    static ov::Shape get_shape(const ggml_tensor * tensor);
    static std::vector<size_t> get_stride(const ggml_tensor * tensor);
    static ov::element::Type get_ov_type(const ggml_tensor * tensor);
    static std::string compute_op_type(const ggml_tensor * node);
    void add_extra_inputs();

    void update_io(ggml_cgraph * cgraph);

    inline static bool is_inp_tok(const ggml_tensor * tensor, const ggml_tensor * op) {
        return op->op == GGML_OP_GET_ROWS && tensor == op->src[1] && op->src[0]->op == GGML_OP_NONE;
    }

    inline static bool is_inp_pos(const ggml_tensor * tensor, const ggml_tensor * op) {
        return op->op == GGML_OP_ROPE && tensor == op->src[1];
    }

    // IMROPE packs 4 stacked position planes (t/h/w/e) into inp_pos, each of length
    // n_tokens; other modes carry a single position per token.
    inline static int get_inp_pos_n_planes(const ggml_tensor * op) {
        return op->op_params[2] == GGML_ROPE_TYPE_IMROPE ? 4 : 1;
    }

    inline static bool is_inp_emb(const ggml_tensor * tensor, const ggml_tensor * op) {
        return tensor->op == GGML_OP_GET_ROWS && op->op == GGML_OP_RMS_NORM;
    }

    inline static bool is_inp_mask(const ggml_tensor * tensor, const ggml_tensor * op) {
        return op->op == GGML_OP_CPY || (op->op == GGML_OP_FLASH_ATTN_EXT && tensor == op->src[3]) ||
               (op->op == GGML_OP_SOFT_MAX && tensor == op->src[1]);
    }

    inline static bool is_rope_freqs_weight(const ggml_tensor * tensor, const ggml_tensor * op) {
        return op->op == GGML_OP_ROPE && tensor == op->src[2];
    }

    // also returns true for cache_s and cache_r in SSM/DeltaNet models
    inline static bool is_kvcache(const ggml_tensor * tensor, const ggml_tensor * op) {
        if (tensor == nullptr) {
            return false;
        }
        return (tensor->buffer != nullptr && tensor->buffer->usage == GGML_BACKEND_BUFFER_USAGE_ANY) ||
               (op != nullptr && op->op == GGML_OP_SET_ROWS && op->src[2] == tensor);
    }

    inline static bool is_kv_idx(const ggml_tensor * tensor, const ggml_tensor * op) {
        return op->op == GGML_OP_SET_ROWS && op->src[1] == tensor;
    }

    inline static bool is_output_idx(const ggml_tensor * tensor, const ggml_tensor * op) {
        return op->op == GGML_OP_GET_ROWS && tensor == op->src[1] && op->src[0]->op != GGML_OP_NONE &&
               op->src[1]->op == GGML_OP_NONE;
    }

    // the state permutation index input used in SSM/DeltaNet models (inp->s_copy in llama-graph.cpp)
    inline static bool is_inp_s_copy(const ggml_tensor * tensor, const ggml_tensor * op) {
        return op->op == GGML_OP_GET_ROWS && tensor == op->src[1] &&
               op->src[0]->buffer->usage == GGML_BACKEND_BUFFER_USAGE_ANY;
    }

    std::string get_graph_input_ov_name(const ggml_tensor * tensor, const ggml_tensor * op) const {
        if (is_inp_pos(tensor, op)) {
            return "inp_pos";
        }
        if (is_inp_emb(tensor, op)) {
            return "embd";
        }
        if (is_stateful() && is_inp_mask(tensor, op)) {
            return std::string(tensor->name).find("swa") == std::string::npos ? "self_kq_mask" : "self_kq_mask_swa";
        }
        return tensor->name;
    }

private:
    void set_input_output();
    int compute_op_case(const ggml_tensor * node) const;
    bool node_is_used_as_src(const int node_idx);
    void compute_model_inputs();
    void compute_model_outputs();

    // True if tensor is the inp->s_copy index leaf gathered by a recurrent state cache GET_ROWS
    // (possibly through a VIEW), so it gets a dynamic [1,1,1,-1] graph-input shape.
    bool is_s_copy_leaf(const ggml_tensor * tensor) const;

    // Infer and propagate dynamic-dimension indices for all tensors in the GGML graph.
    void compute_node_dynamic_dims();

    void validate_cgraph() const;

    ggml_cgraph * m_cgraph = nullptr;
    std::map<std::string, ggml_tensor *> m_inputs;

    std::map<std::string, ov::frontend::ggml::ModelInputInfo> m_model_inputs;
    std::map<std::string, ov::frontend::ggml::ModelExtraInputInfo> m_model_extra_inputs;
    std::map<std::string, std::shared_ptr<ov::Node>> m_model_weights;
    std::map<std::string, ggml_tensor *> m_model_outputs;
    std::set<std::string> m_model_output_names;
    std::vector<NodeInfo> m_node_info_list;
    std::map<ggml_tensor *, int> m_node_dynamic_dims;

    ModelParams m_model_params;
    ComputeParams m_compute_params;
};

void print_tensor_address_map(const ggml_cgraph * cgraph);

std::optional<int> extract_layer_from_name(const std::string & name);