#pragma once

#include "tts_transformer.h"
#include "ggml.h"
#include "ggml-backend.h"
#include "gguf.h"
#include "coreml_code_predictor.h"

#include <map>
#include <string>
#include <vector>
#ifdef QWEN3_TTS_TIMING
#include <chrono>
#endif

namespace qwen3_tts {

#ifdef QWEN3_TTS_TIMING
struct tts_timing {
    double t_prefill_build_ms = 0;
    double t_prefill_forward_ms = 0;
    double t_prefill_graph_build_ms = 0;
    double t_prefill_graph_alloc_ms = 0;
    double t_prefill_compute_ms = 0;
    double t_prefill_data_ms = 0;
    double t_prefill_special_text_proj_ms = 0;
    double t_prefill_instruct_text_proj_ms = 0;
    double t_prefill_role_text_proj_ms = 0;
    double t_prefill_body_text_proj_ms = 0;
    double t_prefill_codec_lookup_ms = 0;
    double t_prefill_ref_code_embed_ms = 0;
    double t_prefill_compose_ms = 0;

    double t_talker_forward_ms = 0;
    double t_talker_graph_build_ms = 0;
    double t_talker_graph_alloc_ms = 0;
    double t_talker_compute_ms = 0;
    double t_talker_data_ms = 0;
    double t_talker_input_upload_ms = 0;
    double t_talker_hidden_read_ms = 0;
    double t_talker_logits_read_ms = 0;
    double t_talker_sched_reset_ms = 0;

    double t_code_pred_ms = 0;
    double t_code_pred_init_ms = 0;
    double t_code_pred_prefill_ms = 0;
    double t_code_pred_steps_ms = 0;
    double t_code_pred_graph_build_ms = 0;
    double t_code_pred_graph_alloc_ms = 0;
    double t_code_pred_compute_ms = 0;
    double t_code_pred_data_ms = 0;
    double t_code_pred_prefill_graph_build_ms = 0;
    double t_code_pred_prefill_graph_alloc_ms = 0;
    double t_code_pred_prefill_compute_ms = 0;
    double t_code_pred_prefill_data_ms = 0;
    double t_code_pred_steps_graph_build_ms = 0;
    double t_code_pred_steps_graph_alloc_ms = 0;
    double t_code_pred_steps_compute_ms = 0;
    double t_code_pred_steps_data_ms = 0;
    double t_code_pred_input_upload_ms = 0;
    double t_code_pred_logits_read_ms = 0;
    double t_code_pred_sampling_ms = 0;
    double t_code_pred_sched_reset_ms = 0;
    double t_code_pred_coreml_ms = 0;

    double t_embed_lookup_ms = 0;

    int32_t n_frames = 0;
    int32_t n_prefill_tokens = 0;
    int32_t n_trailing_tokens = 0;
    double t_generate_total_ms = 0;
};
#endif

#define QWEN3_TTS_MAX_NODES 16384
#define QWEN3_TTS_CODE_PRED_MAX_NODES 4096

struct transformer_layer {
    struct ggml_tensor * attn_norm = nullptr;

    struct ggml_tensor * attn_q = nullptr;
    struct ggml_tensor * attn_k = nullptr;
    struct ggml_tensor * attn_v = nullptr;
    struct ggml_tensor * attn_output = nullptr;
    struct ggml_tensor * attn_q_norm = nullptr;
    struct ggml_tensor * attn_k_norm = nullptr;

    struct ggml_tensor * ffn_norm = nullptr;

    struct ggml_tensor * ffn_gate = nullptr;
    struct ggml_tensor * ffn_up = nullptr;
    struct ggml_tensor * ffn_down = nullptr;
};

struct tts_transformer_model {
    tts_transformer_config config;

    struct ggml_tensor * text_embd = nullptr;
    struct ggml_tensor * text_proj_fc1 = nullptr;
    struct ggml_tensor * text_proj_fc1_bias = nullptr;
    struct ggml_tensor * text_proj_fc2 = nullptr;
    struct ggml_tensor * text_proj_fc2_bias = nullptr;

    struct ggml_tensor * codec_embd = nullptr;

    std::vector<transformer_layer> layers;

    struct ggml_tensor * output_norm = nullptr;
    struct ggml_tensor * codec_head = nullptr;

    std::vector<transformer_layer> code_pred_layers;
    struct ggml_tensor * code_pred_output_norm = nullptr;

    struct ggml_tensor * code_pred_small_to_mtp_weight = nullptr;
    struct ggml_tensor * code_pred_small_to_mtp_bias = nullptr;

    std::vector<struct ggml_tensor *> code_pred_embd;
    std::vector<struct ggml_tensor *> code_pred_head;

    struct ggml_context * ctx = nullptr;
    ggml_backend_buffer_t buffer = nullptr;

    std::map<std::string, struct ggml_tensor *> tensors;
};

struct tts_kv_cache {
    std::vector<struct ggml_tensor *> k_cache;
    std::vector<struct ggml_tensor *> v_cache;

    struct ggml_context * ctx = nullptr;
    ggml_backend_buffer_t buffer = nullptr;

    int32_t n_ctx = 0;
    int32_t n_used = 0;
    int32_t head_dim = 128;
    int32_t n_kv_heads = 8;
    int32_t n_layers = 28;
};

struct tts_transformer_state {
    ggml_backend_t backend = nullptr;
    ggml_backend_t backend_cpu = nullptr;
    ggml_backend_sched_t sched = nullptr;
    ggml_backend_sched_t talker_replay_sched = nullptr;
    ggml_backend_sched_t code_pred_prefill_sched = nullptr;
    ggml_backend_sched_t code_pred_step_sched = nullptr;
    std::vector<ggml_backend_sched_t> code_pred_replay_scheds;
    std::vector<struct ggml_cgraph *> code_pred_replay_graphs;
    bool sched_reserved = false;
    bool sched_reserve_failed = false;
    int32_t sched_reserved_ctx = 0;
    int32_t sched_reserved_prefill_len = 0;
    bool code_pred_sched_reserved = false;
    bool code_pred_sched_reserve_failed = false;
    bool talker_replay_ready = false;
    bool talker_replay_failed = false;
    int32_t talker_replay_n_kv_pad = 0;
    struct ggml_cgraph * talker_replay_graph = nullptr;
    bool code_pred_replay_ready = false;
    bool code_pred_replay_failed = false;

    std::vector<uint8_t> compute_meta;
    std::vector<uint8_t> talker_replay_compute_meta;
    std::vector<std::vector<uint8_t>> code_pred_compute_meta;
    std::vector<ggml_fp16_t> talker_mask;
    std::vector<ggml_fp16_t> code_pred_prefill_mask;
    std::vector<ggml_fp16_t> code_pred_step_masks;
    int32_t code_pred_static_mask_n_ctx = 0;
    struct ggml_tensor * code_pred_prefill_pos = nullptr;
    struct ggml_tensor * code_pred_prefill_mask_tensor = nullptr;
    std::vector<struct ggml_tensor *> code_pred_step_pos;
    std::vector<struct ggml_tensor *> code_pred_step_mask_tensors;
    std::vector<uint8_t> code_pred_graph_stats_logged;

    struct ggml_context * hidden_bridge_ctx = nullptr;
    ggml_backend_buffer_t hidden_bridge_buffer = nullptr;
    struct ggml_tensor * hidden_bridge = nullptr;

    tts_kv_cache cache;
    tts_kv_cache code_pred_cache;
};

struct tts_transformer_private {
    tts_transformer_model model;
    tts_transformer_state state;
    std::vector<ggml_fp16_t> embd_row_fp16_scratch;
    std::vector<float> cached_special_text_proj;
    std::vector<int32_t> cached_reference_code_key;
    std::vector<float> cached_reference_codec_embed;
    int32_t cached_reference_frames = 0;
    int32_t cached_reference_codebooks = 0;
    CoreMLCodePredictor coreml_code_predictor;
    bool use_coreml_code_predictor = false;
    std::string coreml_code_predictor_path;
    bool skip_ggml_code_pred_layers = false;

#ifdef QWEN3_TTS_TIMING
    tts_timing * timing = nullptr;
#endif
};

void free_transformer_model(tts_transformer_model & model);
void free_tts_kv_cache(tts_kv_cache & cache);

} // namespace qwen3_tts