#ifndef RN_SLOT_MANAGER_H

#define RN_SLOT_MANAGER_H

#include "rn-slot.h"

#include "common.h"

#include "llama.h"

#include <vector>

#include <deque>

#include <map>

#include <functional>

#include <mutex>

#include <thread>

#include <atomic>

#include <condition_variable>

namespace rnllama {

// Forward declarations

struct llama_rn_context;

struct completion_token_output;

// Status structures for exposing slot manager state to JS

struct llama_rn_request_status {

int32_t request_id;

std::string type; // "completion", "embedding", "rerank"

std::string state; // "queued", "processing_prompt", "generating", "done"

size_t prompt_length;

size_t tokens_generated;

double prompt_ms;

double generation_ms;

double tokens_per_second;

};

struct llama_rn_parallel_status {

int32_t n_parallel;

int32_t active_slots;

int32_t queued_requests;

std::vector<llama_rn_request_status> requests;

};

enum class llama_rn_cancel_result {

ACTIVE,

QUEUED,

NOT_FOUND,

};

// Queued request structure

struct llama_rn_queued_request {

int32_t request_id;

llama_rn_slot_task_type task_type;

common_params params;

bool post_sampling_probs = true;

std::vector<llama_token> prompt_tokens;

std::function<void(const completion_token_output&)> on_token;

std::function<void(llama_rn_slot*)> on_complete;

// Media paths for multimodal

std::vector<std::string> media_paths;

std::string prompt_text; // Original prompt text (needed for media processing)

// Chat format parameters

int chat_format;

common_reasoning_format reasoning_format;

std::string generation_prompt;

std::string chat_parser; // Serialized PEG parser for chat output parsing

// Prefill text

std::string prefill_text;

// Embedding parameters

int embd_normalize;

std::function<void(int32_t, const std::vector<float>&)> on_embedding;

// Rerank parameters

std::vector<std::vector<llama_token>> rerank_prompt_tokens;

std::function<void(int32_t, const std::vector<float>&)> on_rerank;

std::shared_ptr<llama_rn_decision_job> decision_job;

// State management

std::string load_state_path; // File path to load state from before processing

std::string save_state_path; // File path to save state to after completion

std::string save_prompt_state_path; // File path to save prompt state to after prompt processing

int32_t load_state_size; // Number of tokens to load (0 or -1 = all tokens)

int32_t save_state_size; // Number of tokens to save (0 or -1 = all tokens)

llama_rn_queued_request() :

request_id(-1),

task_type(SLOT_TASK_TYPE_COMPLETION),

chat_format(0),

reasoning_format(COMMON_REASONING_FORMAT_NONE),

embd_normalize(-1),

load_state_size(-1),

save_state_size(-1)

{}

};

// Slot manager for parallel decoding

struct llama_rn_slot_manager {

// Parent context reference

llama_rn_context* parent_ctx;

// Slots own non-copyable batches; grow the pool without relocating them.

std::deque<llama_rn_slot> slots;

int32_t n_parallel; // Number of parallel slots

// Request queue

std::deque<llama_rn_queued_request> queue_requests;

// Request tracking

std::map<int32_t, llama_rn_slot*> active_requests; // request_id -> slot

std::atomic<int32_t> next_request_id;

// Batch processing

llama_batch batch;

int32_t n_batch; // Max batch size

// Shared MTP speculative decoding state. llama.cpp's MTP driver is

// multi-sequence, so queued slots borrow this instead of creating one

// speculative context per slot.

common_speculative *mtp_spec = nullptr;

llama_context *mtp_spec_ctx = nullptr;

int32_t mtp_spec_n_max = 0;

int32_t mtp_spec_n_min = 0;

float mtp_spec_p_min = 0.0f;

bool mtp_spec_backend_sampling = true;

int32_t mtp_spec_n_gpu_layers = -1;

ggml_type mtp_spec_cache_type_k = GGML_TYPE_F16;

ggml_type mtp_spec_cache_type_v = GGML_TYPE_F16;

// Configuration

float slot_prompt_similarity; // Threshold for cache reuse (0.0-1.0)

bool continuous_batching; // Allow mixing prompt/generation

// Processing loop control

std::mutex slots_mutex; // Mutex for thread-safe access to slots

std::condition_variable slots_cv; // Condition variable for efficient waiting

std::thread processing_thread; // Background processing thread

std::atomic<bool> processing_active; // Flag to control processing loop

// Status subscription support

std::map<int32_t, std::function<void(const llama_rn_parallel_status&)>> status_subscribers;

std::mutex subscribers_mutex;

int32_t next_subscriber_id = 1;

// Constructor

llama_rn_slot_manager(llama_rn_context* ctx);

// Destructor

~llama_rn_slot_manager();

// Initialization

bool init(int32_t n_parallel, int32_t n_batch, int32_t n_ctx);

// Request management

int32_t reserve_request_id();

int32_t queue_request(

const common_params& params,

const std::vector<llama_token>& prompt,

const std::vector<std::string>& media_paths,

const std::string& prompt_text,

int chat_format,

common_reasoning_format reasoning_format,

const std::string& generation_prompt,

const std::string& chat_parser,

const std::string& prefill_text,

const std::string& load_state_path,

const std::string& save_state_path,

const std::string& save_prompt_state_path,

int32_t load_state_size,

int32_t save_state_size,

std::function<void(const completion_token_output&)> on_token,

std::function<void(llama_rn_slot*)> on_complete,

int32_t request_id = -1

);

int32_t queue_embedding_request(

const std::vector<llama_token>& tokens,

int embd_normalize,

std::function<void(int32_t, const std::vector<float>&)> on_result,

int32_t request_id = -1

);

int32_t queue_rerank_request(

const std::string& query,

const std::vector<std::string>& documents,

int normalize,

std::function<void(int32_t, const std::vector<float>&)> on_results,

int32_t request_id = -1

);

// Throws std::invalid_argument / std::runtime_error if the request cannot be answered;

// a failure after it is queued goes to on_result as {"error": message}

int32_t queue_decision_request(

const json& body,

std::function<void(int32_t, const json&)> on_result,

int32_t request_id = -1

);

// Slot management

llama_rn_slot* get_available_slot(const std::vector<llama_token>& prompt);

llama_rn_slot* get_slot_by_request_id(int32_t request_id);

void release_slot(llama_rn_slot* slot);

llama_rn_cancel_result cancel_request(int32_t request_id);

// Processing loop management

void start_processing_loop();

void stop_processing_loop();

// Main processing loop (protected by mutex)

void update_slots();

// Helper methods

float compute_similarity(const std::vector<llama_token>& a,

const std::vector<llama_token>& b);

void build_batch();

bool process_batch();

void sample_and_callback();

// Finish a slot's generation: flush the UTF-8 gate, mark done, notify

void complete_slot(llama_rn_slot & slot);

common_speculative* ensure_mtp_speculative(common_params& params);

llama_context* get_mtp_draft_context() const;

void reset_mtp_speculative();

// Process pending queue

void process_pending_queue();

// Evaluate the next prompt of every slot's decision

void process_decision_slots();

void reset_decision_sequence(llama_rn_slot & slot);

// Release completed slots

void release_completed_slots();

// Status methods

llama_rn_parallel_status get_status();

bool has_pending_work();

void notify_status_change();

int32_t add_status_subscriber(std::function<void(const llama_rn_parallel_status&)> callback);

void remove_status_subscriber(int32_t subscriber_id);

};

} // namespace rnllama

#endif /* RN_SLOT_MANAGER_H */