#include "rn-slot-manager.h"
#include "rn-completion.h"
#include "rn-llama.h"
#include "rn-mtmd.hpp"
#include "rn-common.hpp"
#include "ggml.h"
#include <algorithm>
#include <chrono>
#include <cstring>
#include <stdexcept>

namespace rnllama {

// Constructor
llama_rn_slot_manager::llama_rn_slot_manager(llama_rn_context* ctx) :
    parent_ctx(ctx),
    n_parallel(1),
    next_request_id(1),
    n_batch(512),
    slot_prompt_similarity(0.5f),
    continuous_batching(false),
    processing_active(false)
{
    // Initialize batch to zero/null - will be properly allocated later
    std::memset(&batch, 0, sizeof(batch));
}

// Destructor
llama_rn_slot_manager::~llama_rn_slot_manager() {
    // Stop processing loop if active
    stop_processing_loop();

    reset_mtp_speculative();

    // Free batch
    if (batch.token != nullptr) {
        llama_batch_free(batch);
    }

    // Slots will be freed automatically by vector destructor
}

// Initialize slot manager
bool llama_rn_slot_manager::init(int32_t n_parallel_, int32_t n_batch_, int32_t n_ctx) {
    n_parallel = n_parallel_;
    n_batch = n_batch_;

    LOG_INFO("Initializing slot manager with %d parallel slots, batch size %d", n_parallel, n_batch);

    // Allocate slots
    slots.resize(n_parallel);

    // Initialize each slot
    int32_t n_ctx_per_slot = n_ctx / n_parallel;
    for (int32_t i = 0; i < n_parallel; i++) {
        slots[i].id = i;
        slots[i].parent_ctx = parent_ctx;
        slots[i].n_ctx = n_ctx_per_slot;
        slots[i].state = SLOT_STATE_IDLE;
        slots[i].request_id = -1;
        LOG_VERBOSE("Slot %d initialized with context size %d", i, n_ctx_per_slot);
    }

    // Allocate batch
    batch = llama_batch_init(n_batch, 0, n_parallel);
    if (batch.token == nullptr) {
        LOG_ERROR("Failed to allocate batch");
        return false;
    }

    LOG_INFO("Slot manager initialized successfully");
    return true;
}

common_speculative* llama_rn_slot_manager::ensure_mtp_speculative(common_params& params) {
    if (parent_ctx == nullptr || parent_ctx->ctx == nullptr) {
        throw std::runtime_error("MTP speculative decoding requires an initialized context");
    }

    const auto& draft = params.speculative.draft;
    const bool compatible =
        mtp_spec != nullptr &&
        mtp_spec_ctx != nullptr &&
        mtp_spec_n_max == draft.n_max &&
        mtp_spec_n_min == draft.n_min &&
        mtp_spec_p_min == draft.p_min &&
        mtp_spec_backend_sampling == draft.backend_sampling &&
        mtp_spec_n_gpu_layers == draft.n_gpu_layers &&
        mtp_spec_cache_type_k == draft.cache_type_k &&
        mtp_spec_cache_type_v == draft.cache_type_v;

    if (compatible) {
        params.speculative.draft.ctx_tgt = parent_ctx->ctx;
        params.speculative.draft.ctx_dft = mtp_spec_ctx;
        return mtp_spec;
    }

    if (mtp_spec != nullptr || mtp_spec_ctx != nullptr) {
        for (const auto& slot : slots) {
            if (slot.spec_is_shared && slot.state == SLOT_STATE_GENERATING) {
                throw std::runtime_error("cannot change MTP speculative parameters while queued MTP slots are active");
            }
        }
        reset_mtp_speculative();
    }

    mtp_spec_ctx = parent_ctx->createMTPDraftContext(params);
    if (mtp_spec_ctx == nullptr) {
        throw std::runtime_error("failed to create MTP draft context");
    }

    params.speculative.draft.ctx_tgt = parent_ctx->ctx;
    params.speculative.draft.ctx_dft = mtp_spec_ctx;

    const uint32_t n_seq = std::max<int32_t>(1, n_parallel);
    mtp_spec = common_speculative_init(params.speculative, n_seq);
    if (mtp_spec == nullptr) {
        llama_free(mtp_spec_ctx);
        mtp_spec_ctx = nullptr;
        throw std::runtime_error("failed to initialize MTP speculative decoding");
    }

    mtp_spec_n_max = draft.n_max;
    mtp_spec_n_min = draft.n_min;
    mtp_spec_p_min = draft.p_min;
    mtp_spec_backend_sampling = draft.backend_sampling;
    mtp_spec_n_gpu_layers = draft.n_gpu_layers;
    mtp_spec_cache_type_k = draft.cache_type_k;
    mtp_spec_cache_type_v = draft.cache_type_v;

    LOG_INFO("Initialized shared MTP speculative state for %d queued slots", n_parallel);
    return mtp_spec;
}

llama_context* llama_rn_slot_manager::get_mtp_draft_context() const {
    return mtp_spec_ctx;
}

void llama_rn_slot_manager::reset_mtp_speculative() {
    for (auto& slot : slots) {
        if (!slot.spec_is_shared) {
            continue;
        }
        if (slot.spec == mtp_spec) {
            slot.spec = nullptr;
        }
        if (slot.spec_ctx == mtp_spec_ctx) {
            slot.spec_ctx = nullptr;
        }
        slot.spec_is_shared = false;
    }

    if (mtp_spec != nullptr) {
        common_speculative_free(mtp_spec);
        mtp_spec = nullptr;
    }
    if (mtp_spec_ctx != nullptr) {
        llama_free(mtp_spec_ctx);
        mtp_spec_ctx = nullptr;
    }

    mtp_spec_n_max = 0;
    mtp_spec_n_min = 0;
    mtp_spec_p_min = 0.0f;
    mtp_spec_backend_sampling = true;
    mtp_spec_n_gpu_layers = -1;
    mtp_spec_cache_type_k = LM_GGML_TYPE_F16;
    mtp_spec_cache_type_v = LM_GGML_TYPE_F16;
}

int32_t llama_rn_slot_manager::reserve_request_id() {
    return next_request_id.fetch_add(1, std::memory_order_relaxed);
}

// Queue a new request
int32_t llama_rn_slot_manager::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
) {
    if (request_id == -1) {
        request_id = reserve_request_id();
    }

    LOG_INFO("Queuing request %d with %zu prompt tokens (load_state=%s, save_state=%s, save_prompt_state=%s, load_size=%d, save_size=%d)",
             request_id, prompt.size(),
             load_state_path.empty() ? "no" : load_state_path.c_str(),
             save_state_path.empty() ? "no" : save_state_path.c_str(),
             save_prompt_state_path.empty() ? "no" : save_prompt_state_path.c_str(),
             load_state_size,
             save_state_size);

    // Create queued request
    llama_rn_queued_request request;
    request.request_id = request_id;
    request.task_type = SLOT_TASK_TYPE_COMPLETION;
    request.params = params;
    request.prompt_tokens = prompt;
    request.media_paths = media_paths;
    request.prompt_text = prompt_text;
    request.chat_format = chat_format;
    request.reasoning_format = reasoning_format;
    request.generation_prompt = generation_prompt;
    request.chat_parser = chat_parser;
    request.prefill_text = utf8_sanitize(prefill_text);
    request.load_state_path = load_state_path;
    request.save_state_path = save_state_path;
    request.save_prompt_state_path = save_prompt_state_path;
    request.load_state_size = load_state_size;
    request.save_state_size = save_state_size;
    request.on_token = on_token;
    request.on_complete = on_complete;

    // Add to queue
    {
        std::lock_guard<std::mutex> lock(slots_mutex);
        queue_requests.emplace_back(std::move(request));
    }

    // Notify processing thread that new work is available
    slots_cv.notify_one();

    // Notify subscribers of status change (new request queued)
    bool has_subscribers = false;
    {
        std::lock_guard<std::mutex> lock(subscribers_mutex);
        has_subscribers = !status_subscribers.empty();
    }
    if (has_subscribers) {
        notify_status_change();
    }

    return request_id;
}

// Queue an embedding task for parallel processing
int32_t llama_rn_slot_manager::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
) {
    if (parent_ctx == nullptr || parent_ctx->model == nullptr || parent_ctx->ctx == nullptr) {
        LOG_ERROR("Cannot queue embedding: context not initialized");
        return -1;
    }

    if (request_id == -1) {
        request_id = reserve_request_id();
    }

    if (!parent_ctx->params.embedding) {
        LOG_WARNING("Embedding disabled in model parameters; returning zero vector");
        if (on_result) {
            const int n_embd = llama_model_n_embd(parent_ctx->model);
            std::vector<float> empty_embedding(n_embd, 0.0f);
            on_result(request_id, empty_embedding);
        }
        return request_id;
    }

    llama_rn_queued_request request;
    request.request_id = request_id;
    request.task_type = SLOT_TASK_TYPE_EMBEDDING;
    request.prompt_tokens = tokens;
    request.embd_normalize = embd_normalize;
    request.on_embedding = on_result;

    {
        std::lock_guard<std::mutex> lock(slots_mutex);
        queue_requests.emplace_back(std::move(request));
    }

    slots_cv.notify_one();

    // Notify subscribers of status change (new request queued)
    bool has_subscribers = false;
    {
        std::lock_guard<std::mutex> lock(subscribers_mutex);
        has_subscribers = !status_subscribers.empty();
    }
    if (has_subscribers) {
        notify_status_change();
    }

    return request_id;
}

// Queue a rerank task for parallel processing
int32_t llama_rn_slot_manager::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
) {
    if (parent_ctx == nullptr || parent_ctx->model == nullptr || parent_ctx->ctx == nullptr) {
        LOG_ERROR("Cannot queue rerank: context not initialized");
        return -1;
    }

    if (request_id == -1) {
        request_id = reserve_request_id();
    }

    const enum llama_pooling_type pooling_type = llama_pooling_type(parent_ctx->ctx);
    if (pooling_type != LLAMA_POOLING_TYPE_RANK) {
        LOG_ERROR("Reranking not supported by current model (pooling_type=%d)", pooling_type);
        if (on_results) {
            std::vector<float> scores(documents.size(), -1e6f);
            on_results(request_id, scores);
        }
        return request_id;
    }

    if (!parent_ctx->params.embedding) {
        LOG_ERROR("Embedding disabled but required for reranking");
        if (on_results) {
            std::vector<float> scores(documents.size(), -1e6f);
            on_results(request_id, scores);
        }
        return request_id;
    }

    const llama_vocab* vocab = llama_model_get_vocab(parent_ctx->model);
    if (vocab == nullptr) {
        LOG_ERROR("Failed to get vocabulary for rerank task");
        if (on_results) {
            std::vector<float> scores(documents.size(), -1e6f);
            on_results(request_id, scores);
        }
        return request_id;
    }

    llama_rn_queued_request request;
    request.request_id = request_id;
    request.task_type = SLOT_TASK_TYPE_RERANK;
    request.embd_normalize = normalize;
    request.on_rerank = on_results;

    try {
        std::vector<llama_token> query_tokens = common_tokenize(vocab, query, false, true);
        request.rerank_prompt_tokens.reserve(documents.size());

        const bool add_bos = llama_vocab_get_add_bos(vocab);
        const bool is_enc_dec = llama_model_has_encoder(parent_ctx->model);

        for (const std::string& doc : documents) {
            std::vector<llama_token> doc_tokens = common_tokenize(vocab, doc, false, true);
            std::vector<llama_token> rerank_tokens = format_rerank_tokens(vocab, query_tokens, doc_tokens);

            // Convert tokens back to text and re-tokenize using context-aware settings
            std::string rerank_text = tokens_to_str(parent_ctx->ctx, rerank_tokens.begin(), rerank_tokens.end());
            std::vector<llama_token> prompt_tokens = common_tokenize(
                parent_ctx->ctx,
                rerank_text,
                add_bos || is_enc_dec,
                true
            );

            request.rerank_prompt_tokens.push_back(std::move(prompt_tokens));
        }

        if (!request.rerank_prompt_tokens.empty()) {
            request.prompt_tokens = request.rerank_prompt_tokens.front();
        }
    } catch (const std::exception& e) {
        LOG_ERROR("Failed to tokenize rerank inputs: %s", e.what());
        if (on_results) {
            std::vector<float> scores(documents.size(), -1e6f);
            on_results(request_id, scores);
        }
        return request_id;
    }

    if (request.rerank_prompt_tokens.empty()) {
        LOG_INFO("Rerank request %d has no documents; returning empty result", request_id);
        if (on_results) {
            on_results(request_id, {});
        }
        return request_id;
    }

    {
        std::lock_guard<std::mutex> lock(slots_mutex);
        queue_requests.emplace_back(std::move(request));
    }

    slots_cv.notify_one();

    // Notify subscribers of status change (new request queued)
    bool has_subscribers = false;
    {
        std::lock_guard<std::mutex> lock(subscribers_mutex);
        has_subscribers = !status_subscribers.empty();
    }
    if (has_subscribers) {
        notify_status_change();
    }

    return request_id;
}

// Get available slot (LRU strategy for now, similarity matching in Phase 3)
llama_rn_slot* llama_rn_slot_manager::get_available_slot(const std::vector<llama_token>& prompt) {
    llama_rn_slot* best_slot = nullptr;
    int64_t oldest_time = INT64_MAX;

    // Find idle or done slot with oldest t_last_used (LRU)
    for (auto& slot : slots) {
        if (slot.state == SLOT_STATE_IDLE || slot.state == SLOT_STATE_DONE) {
            if (slot.t_last_used < oldest_time) {
                oldest_time = slot.t_last_used;
                best_slot = &slot;
            }
        }
    }

    if (best_slot != nullptr) {
        LOG_VERBOSE("Selected slot %d (LRU)", best_slot->id);
    } else {
        LOG_VERBOSE("No available slots");
    }

    return best_slot;
}

// Get slot by request ID
llama_rn_slot* llama_rn_slot_manager::get_slot_by_request_id(int32_t request_id) {
    auto it = active_requests.find(request_id);
    if (it != active_requests.end()) {
        return it->second;
    }
    return nullptr;
}

// Release slot
void llama_rn_slot_manager::release_slot(llama_rn_slot* slot) {
    LOG_VERBOSE("Releasing slot %d", slot->id);

    // Update last used timestamp for LRU tracking
    slot->t_last_used = lm_ggml_time_us();

    // Reset slot (cache_tokens is preserved by reset() for potential reuse)
    slot->reset();
}

// Cancel request
llama_rn_cancel_result llama_rn_slot_manager::cancel_request(int32_t request_id) {
    LOG_INFO("Cancelling request %d", request_id);

    llama_rn_cancel_result result = llama_rn_cancel_result::NOT_FOUND;
    {
        std::lock_guard<std::mutex> lock(slots_mutex);

        auto active_it = active_requests.find(request_id);
        if (active_it != active_requests.end() &&
            (active_it->second->state == SLOT_STATE_PROCESSING_PROMPT ||
             active_it->second->state == SLOT_STATE_GENERATING)) {
            // The processing thread owns terminalization and slot release.
            active_it->second->is_interrupted = true;
            LOG_INFO(
                "Request %d cancellation requested (active in slot %d)",
                request_id,
                active_it->second->id
            );
            result = llama_rn_cancel_result::ACTIVE;
        } else {
            auto queued_it = std::find_if(
                queue_requests.begin(),
                queue_requests.end(),
                [request_id](const llama_rn_queued_request& request) {
                    return request.request_id == request_id;
                }
            );
            if (queued_it != queue_requests.end()) {
                queue_requests.erase(queued_it);
                LOG_INFO("Request %d cancelled (was in pending queue)", request_id);
                result = llama_rn_cancel_result::QUEUED;
            }
        }
    }

    if (result == llama_rn_cancel_result::NOT_FOUND) {
        LOG_WARNING("Request %d not found for cancellation", request_id);
        return result;
    }

    // Wake the processing thread after releasing slots_mutex. Active requests
    // still need worker-owned completion; queued removals may leave it idle.
    slots_cv.notify_one();

    // Notify subscribers of status change
    bool has_subscribers = false;
    {
        std::lock_guard<std::mutex> lock(subscribers_mutex);
        has_subscribers = !status_subscribers.empty();
    }
    if (has_subscribers) {
        notify_status_change();
    }

    return result;
}

// Compute similarity between two token sequences (stub for Phase 3)
float llama_rn_slot_manager::compute_similarity(
    const std::vector<llama_token>& a,
    const std::vector<llama_token>& b
) {
    // Longest Common Prefix (LCP) approach
    size_t common_prefix = 0;
    size_t max_len = std::min(a.size(), b.size());
    for (size_t i = 0; i < max_len; i++) {
        if (a[i] == b[i]) {
            common_prefix++;
        } else {
            break;
        }
    }

    if (a.empty() && b.empty()) return 1.0f;
    if (a.empty() || b.empty()) return 0.0f;

    return static_cast<float>(common_prefix) / static_cast<float>(std::max(a.size(), b.size()));
}

// Process pending queue
void llama_rn_slot_manager::process_pending_queue() {
    while (!queue_requests.empty()) {
        llama_rn_queued_request& request = queue_requests.front();

        const std::vector<llama_token>* prompt_view = nullptr;
        std::vector<llama_token> empty_prompt;

        if (request.task_type == SLOT_TASK_TYPE_RERANK) {
            if (!request.rerank_prompt_tokens.empty()) {
                prompt_view = &request.rerank_prompt_tokens.front();
            }
        } else {
            prompt_view = &request.prompt_tokens;
        }

        if (prompt_view == nullptr) {
            prompt_view = &empty_prompt;
        }

        llama_rn_slot* slot = get_available_slot(*prompt_view);
        if (slot == nullptr) {
            LOG_VERBOSE(
                "No available slots, stopping queue processing (request %d at front)",
                request.request_id
            );
            break;
        }

        // Assign request to slot
        // A cancelled slot can be reassigned before release_slot() has reset
        // it; generation state must not leak into the new request
        slot->clear_generation_state();
        slot->request_id = request.request_id;
        slot->task_type = request.task_type;
        slot->is_interrupted = false;

        // Reset callbacks from previous usage
        slot->on_token_callback = nullptr;
        slot->on_complete_callback = nullptr;
        slot->on_embedding_callback = nullptr;
        slot->on_rerank_callback = nullptr;

        // Ensure we start without a sampling context unless set below
        if (slot->ctx_sampling != nullptr) {
            common_sampler_free(slot->ctx_sampling);
            slot->params = nullptr;
            slot->ctx_sampling = nullptr;
        }

        switch (request.task_type) {
            case SLOT_TASK_TYPE_COMPLETION: {
                slot->params_storage = request.params;
                slot->params = &slot->params_storage;
                slot->ctx_sampling = common_sampler_init(parent_ctx->model, slot->params->sampling);

                // Assign state parameters
                slot->load_state_path = request.load_state_path;
                slot->save_state_path = request.save_state_path;
                slot->save_prompt_state_path = request.save_prompt_state_path;
                slot->load_state_size = request.load_state_size;
                slot->save_state_size = request.save_state_size;

                // Load state if provided
                if (!slot->load_state_path.empty()) {
                    if (!slot->load_state()) {
                        LOG_ERROR("Failed to load state for slot %d, request %d",
                                  slot->id, request.request_id);
                        // Mark slot as done with error
                        slot->state = SLOT_STATE_DONE;
                        slot->incomplete = true;
                        slot->error_message = "Failed to load state from: " + slot->load_state_path;
                        if (request.on_complete) {
                            request.on_complete(slot);
                        }
                        queue_requests.pop_front();
                        continue;
                    }
                }

                // Start timing AFTER state loading completes
                slot->t_start_process = lm_ggml_time_us();

                // Always load prompt - it will detect and preserve state if appropriate
                bool has_media = !request.media_paths.empty();
                if (has_media && parent_ctx->isMultimodalEnabled()) {
                    LOG_INFO("Storing %zu media paths for deferred processing in slot %d",
                             request.media_paths.size(), slot->id);
                    slot->media_paths = request.media_paths;
                    slot->prompt_text = request.prompt_text;
                    slot->media_processed = false;
                    slot->load_prompt(request.prompt_tokens);
                } else {
                    slot->media_paths.clear();
                    slot->prompt_text.clear();
                    slot->media_processed = true;
                    slot->load_prompt(request.prompt_tokens);
                }
                slot->i_batch = -1;

                slot->on_token_callback = request.on_token;
                slot->on_complete_callback = request.on_complete;
                slot->current_chat_format = request.chat_format;
                slot->current_reasoning_format = request.reasoning_format;
                slot->current_generation_prompt = request.generation_prompt;
                slot->current_chat_parser = request.chat_parser;
                slot->prefill_text = request.prefill_text;
                slot->n_remaining = request.params.n_predict;
                slot->stop_words = request.params.antiprompt;
                break;
            }

            case SLOT_TASK_TYPE_EMBEDDING: {
                slot->params_storage = request.params;
                slot->params = &slot->params_storage;
                // Start timing (no state loading for embeddings)
                slot->t_start_process = lm_ggml_time_us();

                slot->media_paths.clear();
                slot->prompt_text.clear();
                slot->media_processed = true;
                slot->embd_normalize = request.embd_normalize;
                slot->on_embedding_callback = request.on_embedding;
                slot->n_remaining = -1;
                slot->stop_words.clear();
                slot->load_prompt(request.prompt_tokens);
                slot->i_batch = -1;
                break;
            }

            case SLOT_TASK_TYPE_RERANK: {
                slot->params = nullptr;
                // Start timing (memory clear is part of the task, not overhead)
                slot->t_start_process = lm_ggml_time_us();

                if (parent_ctx && parent_ctx->ctx) {
                    // Only this slot's sequence - a global clear would corrupt
                    // other slots' in-flight sequences
                    llama_memory_seq_rm(llama_get_memory(parent_ctx->ctx), slot->id, 0, -1);
                }
                if (request.rerank_prompt_tokens.empty()) {
                    LOG_WARNING("Rerank request %d has no documents to process", request.request_id);
                    if (request.on_rerank) {
                        request.on_rerank(request.request_id, {});
                    }
                    queue_requests.pop_front();
                    continue;
                }

                slot->media_paths.clear();
                slot->prompt_text.clear();
                slot->media_processed = true;
                slot->embd_normalize = request.embd_normalize;
                slot->on_rerank_callback = request.on_rerank;
                slot->rerank_prompt_tokens = std::move(request.rerank_prompt_tokens);
                slot->rerank_scores.assign(slot->rerank_prompt_tokens.size(), 0.0f);
                slot->rerank_current_index = 0;
                slot->n_remaining = -1;
                slot->stop_words.clear();
                slot->load_prompt(slot->rerank_prompt_tokens[0]);
                slot->i_batch = -1;
                break;
            }

            default:
                LOG_ERROR("Unknown task type %d for request %d", request.task_type, request.request_id);
                queue_requests.pop_front();
                continue;
        }

        // Track active request
        active_requests[request.request_id] = slot;

        // Remove from queue
        queue_requests.pop_front();
    }
}

// Build batch from all active slots
void llama_rn_slot_manager::build_batch() {
    // Clear the batch
    batch.n_tokens = 0;

    // First pass: Add tokens from GENERATING slots (previously sampled tokens)
    for (auto& slot : slots) {
        if (slot.state == SLOT_STATE_GENERATING) {
            if (slot.task_type == SLOT_TASK_TYPE_COMPLETION && slot.should_use_mtp()) {
                continue;
            }
            // Only add if we have generated tokens (skip first iteration after prompt)
            if (!slot.generated_tokens.empty()) {
                // Get the last generated token
                llama_token token = slot.generated_tokens.back();

                // Add to batch with this slot's sequence ID
                llama_batch_add(&batch, token, slot.n_past, {slot.id}, true);

                // Mark position in batch for this slot
                slot.i_batch = batch.n_tokens - 1;

                slot.n_past++; // Increment for next token

                LOG_VERBOSE("Slot %d: Added generated token %d at pos %d", slot.id, token, slot.n_past - 1);
            }
        }
    }

    // Second pass: Add prompt tokens from PROCESSING_PROMPT slots
    for (auto& slot : slots) {
        if (slot.state == SLOT_STATE_PROCESSING_PROMPT) {
            if (slot.task_type == SLOT_TASK_TYPE_COMPLETION && slot.should_use_mtp()) {
                if (!slot.media_paths.empty()) {
                    LOG_ERROR("Slot %d: MTP speculative decoding does not support media inputs", slot.id);
                    slot.incomplete = true;
                    slot.error_message = "MTP speculative decoding currently supports text-only queued completions";
                    complete_slot(slot);
                    continue;
                }

                slot.state = SLOT_STATE_GENERATING;
                slot.i_batch = -1;
                LOG_INFO("Slot %d: Transitioned to GENERATING state with MTP speculative decoding", slot.id);
                continue;
            }

            // Check if we need to process media first (deferred processing)
            if (!slot.media_processed && !slot.media_paths.empty()) {
                LOG_INFO("Slot %d: Processing media before prompt tokens", slot.id);

                try {
                    // Seed the media evaluator with the cached history (loaded
                    // state or a previous turn on this slot) so it can reuse
                    // the sequence memory; processMedia reconciles or clears
                    // the memory itself. The history is only trustworthy where
                    // the memory backs it.
                    slot.embd = slot.cache_tokens;
                    if (parent_ctx && parent_ctx->ctx) {
                        auto * kv = llama_get_memory(parent_ctx->ctx);
                        const llama_pos mem_len = llama_memory_seq_pos_max(kv, slot.id) + 1;
                        // M-RoPE media histories legitimately hold fewer time
                        // positions than placeholder tokens - leave them as-is
                        const bool mrope_media =
                            model_uses_mrope(llama_get_model(parent_ctx->ctx)) &&
                            std::find(slot.embd.begin(), slot.embd.end(),
                                      LLAMA_TOKEN_NULL) != slot.embd.end();
                        if ((llama_pos) slot.embd.size() > mem_len && !mrope_media) {
                            // The last sampled token of a previous run may
                            // never have been decoded; keep the decoded prefix
                            slot.embd.resize(mem_len);
                        } else if ((llama_pos) slot.embd.size() < mem_len) {
                            // Memory holds positions the history doesn't
                            // describe (e.g. cleared elsewhere) - not reusable
                            LOG_WARNING("Slot %d: Cached history (%zu tokens) does not match memory (%lld positions), discarding",
                                       slot.id, slot.embd.size(), (long long) mem_len);
                            slot.embd.clear();
                            slot.bitmap_past_hashes.clear();
                            llama_memory_seq_rm(kv, slot.id, 0, -1);
                        }
                    }
                    bool context_full = false;

                    // Recurrent/hybrid prompt checkpoints must be written at a
                    // chunk-aligned anchor while the memory holds exactly those
                    // positions: a post-eval save could never be rolled back to
                    // a resumable point on reload. processMedia fires capture
                    // at the reused frontier and right after the last media
                    // chunk; the last write (the media boundary) wins, and the
                    // short text tail is simply re-decoded on reload.
                    bool prompt_ckpt_written = false;
                    mtmd_state_capture_fn capture = nullptr;
                    const llama_model * mdl = llama_get_model(parent_ctx->ctx);
                    if ((llama_model_is_recurrent(mdl) || llama_model_is_hybrid(mdl)) &&
                        slot.save_prompt_state_pending && !slot.save_prompt_state_path.empty()) {
                        // The capture overwrites the state file mid-eval; drop
                        // any previous sidecar now so an eval failure (or kill)
                        // before the new sidecar is written can never pair
                        // stale hashes with the new file (fail closed)
                        write_state_meta(slot.save_prompt_state_path, {});
                        auto * slot_ptr = &slot;
                        capture = [this, slot_ptr, &prompt_ckpt_written](
                                      const std::vector<llama_token> &toks, size_t n) {
                            if (n == 0 || n > toks.size()) {
                                return;
                            }
                            std::vector<llama_token> prefix(toks.begin(), toks.begin() + n);
                            const size_t nwrite = llama_state_seq_save_file(
                                parent_ctx->ctx, slot_ptr->save_prompt_state_path.c_str(),
                                slot_ptr->id, prefix.data(), prefix.size());
                            if (nwrite > 0) {
                                prompt_ckpt_written = true;
                                LOG_INFO("Slot %d: Saved media prompt checkpoint at %zu tokens (%.2f KB)",
                                        slot_ptr->id, n, nwrite / 1024.0);
                            } else {
                                LOG_WARNING("Slot %d: Failed to save media prompt checkpoint at %zu tokens",
                                           slot_ptr->id, n);
                            }
                        };
                    }

                    parent_ctx->mtmd_wrapper->processMedia(
                        parent_ctx->ctx,
                        slot.prompt_text,
                        slot.media_paths,
                        parent_ctx->n_ctx,
                        n_batch,
                        slot.n_past,
                        slot.embd,
                        context_full,
                        slot.ctx_sampling,
                        slot.bitmap_past_hashes,
                        slot.id,  // Use slot ID as sequence ID for parallel processing
                        /*recover*/ nullptr,
                        capture,
                        /*invalidate*/ nullptr
                    );

                    if (context_full) {
                        LOG_ERROR("Context full after processing media for slot %d", slot.id);
                        slot.context_full = true;
                        complete_slot(slot);
                        continue;
                    }

                    // Update prompt tokens with the processed result from processMedia
                    slot.prompt_tokens = slot.embd;
                    slot.num_prompt_tokens = slot.embd.size();
                    slot.cache_tokens = slot.embd;
                    slot.media_processed = true;
                    slot.n_prompt_tokens_cache = parent_ctx->mtmd_wrapper->last_reused_n_past;
                    if (prompt_ckpt_written) {
                        // The capture anchor is the resumable file (all media
                        // sits inside it); persist its media identity and skip
                        // the post-eval save
                        write_state_meta(slot.save_prompt_state_path, slot.bitmap_past_hashes);
                        slot.save_prompt_state_pending = false;
                        slot.save_prompt_state_tokens = -1;
                    } else if (slot.save_prompt_state_pending) {
                        // processMedia evaluated the whole prompt already; the
                        // checkpoint token list must cover every position in
                        // the sequence memory to stay resumable
                        slot.save_prompt_state_tokens = (llama_pos)slot.num_prompt_tokens;
                        LOG_VERBOSE("Slot %d: Updated prompt checkpoint target to %lld/%zu tokens after media processing",
                                   slot.id, (long long)slot.save_prompt_state_tokens, slot.num_prompt_tokens);
                    }

                    // A reused prefix can end inside the trailing text (exact
                    // prompt replay / text-tail checkpoint recovery): the chunk
                    // eval loop skips those tokens, so fall through and let the
                    // normal prompt loop below decode them for fresh logits.
                    if ((size_t) slot.n_past < slot.num_prompt_tokens) {
                        LOG_INFO("Slot %d: Media processed, decoding remaining prompt tail from %d/%zu",
                                slot.id, slot.n_past, slot.num_prompt_tokens);
                    } else {
                        // processMedia() evaluated every prompt token, so the
                        // logits of the last one are ready to sample
                        slot.n_past = slot.num_prompt_tokens;

                        // Transition to GENERATING state immediately since all prompt tokens are processed
                        slot.state = SLOT_STATE_GENERATING;

                        // Mark that prompt processing just finished - timing will be calculated after decode
                        // Note: for media processing, processMedia() already decoded everything, so timing
                        // calculation will happen immediately after this in the main loop
                        slot.prompt_processing_finished = true;
                        slot.n_prompt_tokens_processed = slot.num_prompt_tokens - slot.n_prompt_tokens_cache;

                        // Set i_batch to -1 to indicate logits from media processing are ready to sample
                        // In sample_and_callback(), batch index -1 will be handled specially
                        slot.i_batch = -1;

                        // Sample the first token now: the context logits still
                        // belong to processMedia's final decode, and
                        // process_batch() would overwrite them with other
                        // slots' tokens before sample_and_callback runs
                        if (slot.ctx_sampling != nullptr) {
                            slot.media_pending_token =
                                common_sampler_sample(slot.ctx_sampling, parent_ctx->ctx, -1);
                        }

                        LOG_INFO("Slot %d: Media processed, transitioned to GENERATING state, n_past=%d, num_prompt_tokens=%zu",
                                slot.id, slot.n_past, slot.num_prompt_tokens);

                        // Continue to next slot - this slot is ready to sample in sample_and_callback()
                        continue;
                    }

                } catch (const std::exception& e) {
                    LOG_ERROR("Failed to process media for slot %d: %s", slot.id, e.what());
                    // A partial evaluation leaves the sequence memory out of
                    // sync with the cached history - drop both
                    slot.cache_tokens.clear();
                    slot.bitmap_past_hashes.clear();
                    if (parent_ctx && parent_ctx->ctx) {
                        llama_memory_seq_rm(llama_get_memory(parent_ctx->ctx), slot.id, 0, -1);
                    }
                    slot.incomplete = true;
                    complete_slot(slot);
                    continue;
                }
            }

            // Process tokens up to n_batch limit (only for non-media slots)
            size_t prompt_end = slot.num_prompt_tokens;
            if (slot.save_prompt_state_pending && slot.save_prompt_state_tokens >= 0 &&
                slot.n_past <= slot.save_prompt_state_tokens) {
                prompt_end = std::min(prompt_end, (size_t)slot.save_prompt_state_tokens);
            }

            while (slot.n_past < (llama_pos)prompt_end && batch.n_tokens < n_batch) {
                llama_token token = slot.prompt_tokens[slot.n_past];

                // Skip LLAMA_TOKEN_NULL - these are media placeholders already in KV cache
                if (token == LLAMA_TOKEN_NULL) {
                    LOG_VERBOSE("Slot %d: Skipping NULL token at pos %d (media chunk)", slot.id, slot.n_past);
                    slot.n_past++;
                    continue;
                }

                // Request logits for all tokens when embeddings/rerank are needed
                bool need_logits = true;
                if (slot.task_type == SLOT_TASK_TYPE_COMPLETION) {
                    need_logits = (slot.n_past == (llama_pos)(slot.num_prompt_tokens - 1));
                }

                // Add to batch with this slot's sequence ID
                llama_batch_add(&batch, token, slot.n_past, {slot.id}, need_logits);

                // Mark position in batch for this slot (will be overwritten each iteration)
                slot.i_batch = batch.n_tokens - 1;

                slot.n_past++;
            }

            // If we've processed all prompt tokens, transition based on task type
            if (slot.n_past >= (llama_pos)slot.num_prompt_tokens) {
                slot.state = SLOT_STATE_GENERATING;

                // Mark that prompt processing just finished - timing will be calculated after decode
                slot.prompt_processing_finished = true;
                slot.n_prompt_tokens_processed = slot.num_prompt_tokens - slot.n_prompt_tokens_cache;

                if (slot.task_type == SLOT_TASK_TYPE_COMPLETION) {
                    LOG_INFO("Slot %d: Transitioned to GENERATING state", slot.id);
                } else if (slot.task_type == SLOT_TASK_TYPE_EMBEDDING) {
                    LOG_INFO("Slot %d: Prompt processed for embedding task", slot.id);
                } else if (slot.task_type == SLOT_TASK_TYPE_RERANK) {
                    LOG_INFO("Slot %d: Prompt processed for rerank task (doc %zu/%zu)",
                             slot.id,
                             slot.rerank_current_index + 1,
                             slot.rerank_prompt_tokens.size());
                }
            }

            LOG_VERBOSE("Slot %d: Processed prompt tokens, n_past=%d/%zu",
                       slot.id, slot.n_past, slot.num_prompt_tokens);
        }
    }

    LOG_VERBOSE("Batch built with %d tokens", batch.n_tokens);
}

bool llama_rn_slot_manager::process_batch() {
    if (batch.n_tokens == 0) {
        // No tokens to process
        return true;
    }

    if (parent_ctx == nullptr || parent_ctx->ctx == nullptr) {
        LOG_ERROR("Cannot process batch: context is null");
        return false;
    }

    // Call llama_decode with the unified batch
    int ret = llama_decode(parent_ctx->ctx, batch);

    if (ret != 0) {
        // Decode failed
        if (ret == 1) {
            LOG_ERROR("llama_decode failed: could not find a KV slot for the batch");
        } else {
            LOG_ERROR("llama_decode failed with code: %d", ret);
        }

        // Try with smaller batch size next time
        if (n_batch > 32) {
            n_batch = n_batch / 2;
            LOG_WARNING("Reducing batch size to %d", n_batch);
        }

        return false;
    }

    // Synchronize to ensure GPU work completes before timing measurements
    // This is critical for accurate performance metrics when using Metal/GPU
    llama_synchronize(parent_ctx->ctx);

    LOG_VERBOSE("Batch processed successfully");
    return true;
}

void llama_rn_slot_manager::complete_slot(llama_rn_slot & slot) {
    slot.generated_text += slot.utf8_gate.finish();
    slot.state = SLOT_STATE_DONE;
    auto on_complete = std::move(slot.on_complete_callback);
    slot.on_complete_callback = nullptr;
    if (on_complete) {
        on_complete(&slot);
    }
}

void llama_rn_slot_manager::sample_and_callback() {
    if (parent_ctx == nullptr || parent_ctx->ctx == nullptr) {
        return;
    }

    const llama_vocab* vocab = llama_model_get_vocab(parent_ctx->model);
    const int n_embd = llama_model_n_embd(parent_ctx->model);

    auto get_embedding_ptr = [&](llama_rn_slot& slot) -> const float* {
        const float* data = llama_get_embeddings_seq(parent_ctx->ctx, slot.id);
        if (data == nullptr) {
            int idx = slot.i_batch;
            if (idx < 0 || idx >= batch.n_tokens) {
                idx = batch.n_tokens - 1;
            }
            if (idx >= 0) {
                data = llama_get_embeddings_ith(parent_ctx->ctx, idx);
            }
        }
        if (data == nullptr) {
            data = llama_get_embeddings(parent_ctx->ctx);
        }
        return data;
    };

    // Process each slot in GENERATING state
    for (auto& slot : slots) {
        if (slot.state != SLOT_STATE_GENERATING) {
            continue;
        }

        // Check if interrupted
        if (slot.is_interrupted) {
            LOG_INFO("Slot %d: Generation interrupted", slot.id);
            complete_slot(slot);
            continue;
        }
        switch (slot.task_type) {
            case SLOT_TASK_TYPE_COMPLETION: {
                if (slot.ctx_sampling == nullptr) {
                    LOG_WARNING("Slot %d: Sampling context is null, marking as done", slot.id);
                    slot.state = SLOT_STATE_DONE;
                    continue;
                }

                if (slot.should_use_mtp()) {
                    auto finish_slot = [&]() {
                        if (!slot.save_state_path.empty()) {
                            slot.save_state();
                        }
                        complete_slot(slot);
                    };

                    auto emit_token = [&](completion_token_output token_output) -> bool {
                        if (token_output.text.empty()) {
                            token_output.text = common_token_to_piece(parent_ctx->ctx, token_output.tok);
                        }
                        token_output.request_id = slot.request_id;

                        token_output.text = slot.utf8_gate.feed(token_output.text);
                        slot.generated_text += token_output.text;

                        const int64_t t_current = lm_ggml_time_us();
                        slot.t_token_generation = (t_current - slot.t_start_generation) / 1e6;

                        slot.generated_tokens.push_back(token_output.tok);
                        slot.n_decoded++;
                        slot.cache_tokens.push_back(token_output.tok);

                        // still emit an empty delta when it carries requested probs
                        if (slot.on_token_callback && (!token_output.text.empty() || !token_output.probs.empty())) {
                            slot.on_token_callback(token_output);
                        }

                        bool should_stop = false;

                        if (slot.n_remaining > 0) {
                            slot.n_remaining--;
                            if (slot.n_remaining == 0) {
                                slot.stopped_limit = true;
                                should_stop = true;
                                LOG_INFO("Slot %d: Stopped on token limit", slot.id);
                            }
                        }

                        if (slot.context_full) {
                            should_stop = true;
                            LOG_WARNING("Slot %d: Context full", slot.id);
                        }

                        if (!slot.stop_words.empty() && !slot.generated_text.empty()) {
                            const std::string& text = slot.generated_text;
                            const size_t last_token_size = token_output.text.size();

                            for (const std::string& word : slot.stop_words) {
                                const size_t search_start = text.size() > word.size() + last_token_size
                                    ? text.size() - word.size() - last_token_size
                                    : 0;
                                size_t pos = text.find(word, search_start);

                                if (pos != std::string::npos) {
                                    slot.stopped_word = true;
                                    slot.stopping_word = word;
                                    should_stop = true;
                                    LOG_INFO("Slot %d: Stopped on word '%s'", slot.id, word.c_str());
                                    break;
                                }
                            }
                        }

                        return should_stop;
                    };

                    try {
                        bool should_stop = false;

                        completion_token_output token_output = slot.next_token_mtp();
                        const bool emitted = token_output.tok != -1;
                        if (emitted) {
                            should_stop = emit_token(std::move(token_output));
                        }

                        const bool done = should_stop ||
                            slot.stopped_limit ||
                            slot.context_full ||
                            (slot.stopped_eos && slot.spec_pending_tokens.empty());

                        if (done) {
                            finish_slot();
                        } else if (!emitted) {
                            slot.incomplete = true;
                            slot.error_message = "MTP speculative decoding did not produce a token";
                            finish_slot();
                        }
                    } catch (const std::exception& e) {
                        LOG_ERROR("Slot %d: MTP speculative decoding failed: %s", slot.id, e.what());
                        slot.incomplete = true;
                        slot.error_message = e.what();
                        finish_slot();
                    }

                    continue;
                }

                if (slot.i_batch == -1) {
                    LOG_VERBOSE("Slot %d: Sampling from media processing logits (batch index -1)", slot.id);
                } else if (slot.i_batch < 0 || slot.i_batch >= batch.n_tokens) {
                    LOG_WARNING("Slot %d: Invalid batch position %d", slot.id, slot.i_batch);
                    continue;
                }

                llama_token new_token_id;
                if (slot.i_batch == -1 && slot.media_pending_token != LLAMA_TOKEN_NULL) {
                    // Pre-sampled right after media ingest (see build_batch);
                    // the context logits no longer belong to this slot here
                    new_token_id = slot.media_pending_token;
                    slot.media_pending_token = LLAMA_TOKEN_NULL;
                } else {
                    new_token_id = common_sampler_sample(slot.ctx_sampling, parent_ctx->ctx, slot.i_batch);
                }
                common_sampler_accept(slot.ctx_sampling, new_token_id, true);

                if (llama_vocab_is_eog(vocab, new_token_id)) {
                    slot.stopped_eos = true;
                    LOG_INFO("Slot %d: Stopped on EOS token", slot.id);

                    // Save state if path is provided
                    if (!slot.save_state_path.empty()) {
                        slot.save_state();
                    }

                    complete_slot(slot);
                    continue;
                }

                std::string token_text = common_token_to_piece(parent_ctx->ctx, new_token_id);
                token_text = slot.utf8_gate.feed(token_text);
                slot.generated_text += token_text;

                // Update token generation timing
                const int64_t t_current = lm_ggml_time_us();
                slot.t_token_generation = (t_current - slot.t_start_generation) / 1e6;

                completion_token_output token_output;
                token_output.tok = new_token_id;
                token_output.text = token_text;
                token_output.request_id = slot.request_id;

                const int32_t n_probs = slot.params->sampling.n_probs;
                if (n_probs > 0) {
                  llama_token_data_array cur_p = *common_sampler_get_candidates(slot.ctx_sampling, true);
                  for (size_t i = 0; i < std::min(cur_p.size, (size_t)n_probs); ++i)
                  {
                      token_output.probs.push_back({cur_p.data[i].id, cur_p.data[i].p});
                  }
                }

                slot.generated_tokens.push_back(new_token_id);
                slot.n_decoded++;
                slot.num_tokens_predicted++;

                // Update cache_tokens to keep track of all processed tokens
                // This is needed for state saving
                slot.cache_tokens.push_back(new_token_id);

                // still emit an empty delta when it carries requested probs
                if (slot.on_token_callback && (!token_output.text.empty() || !token_output.probs.empty())) {
                    slot.on_token_callback(token_output);
                }

                bool should_stop = false;

                if (slot.n_remaining > 0) {
                    slot.n_remaining--;
                    if (slot.n_remaining == 0) {
                        slot.stopped_limit = true;
                        should_stop = true;
                        LOG_INFO("Slot %d: Stopped on token limit", slot.id);
                    }
                }

                if (slot.n_past >= slot.n_ctx) {
                    slot.context_full = true;
                    should_stop = true;
                    LOG_WARNING("Slot %d: Context full", slot.id);
                }

                if (!slot.stop_words.empty() && !slot.generated_text.empty()) {
                    const std::string& text = slot.generated_text;
                    const size_t last_token_size = token_text.size();

                    for (const std::string& word : slot.stop_words) {
                        const size_t search_start = text.size() > word.size() + last_token_size
                            ? text.size() - word.size() - last_token_size
                            : 0;
                        size_t pos = text.find(word, search_start);

                        if (pos != std::string::npos) {
                            slot.stopped_word = true;
                            slot.stopping_word = word;
                            should_stop = true;
                            LOG_INFO("Slot %d: Stopped on word '%s'", slot.id, word.c_str());
                            break;
                        }
                    }
                }

                if (should_stop) {
                    // Save state if path is provided
                    if (!slot.save_state_path.empty()) {
                        slot.save_state();
                    }

                    complete_slot(slot);
                }

                LOG_VERBOSE("Slot %d: Generated token %d ('%s'), n_past=%d, n_decoded=%d",
                           slot.id, new_token_id, token_text.c_str(), slot.n_past, slot.n_decoded);
                break;
            }

            case SLOT_TASK_TYPE_EMBEDDING: {
                const float* data = get_embedding_ptr(slot);

                std::vector<float> embedding(n_embd, 0.0f);
                if (data != nullptr) {
                    embedding.assign(data, data + n_embd);
                }

                std::vector<float> normalized(n_embd, 0.0f);
                if (n_embd >= 4) {
                    LOG_INFO("Embedding data: 0: %f 1: %f 2: %f 3: %f",
                             embedding[0], embedding[1], embedding[2], embedding[3]);
                }
                LOG_INFO("Normalizing embedding with normalize=%d", slot.embd_normalize);
                common_embd_normalize(embedding.data(), normalized.data(), n_embd, slot.embd_normalize);
                if (n_embd >= 4) {
                    LOG_INFO("Normalized embedding data: 0: %f 1: %f 2: %f 3: %f",
                             normalized[0], normalized[1], normalized[2], normalized[3]);
                }

                if (slot.on_embedding_callback) {
                    slot.on_embedding_callback(slot.request_id, normalized);
                }

                slot.state = SLOT_STATE_DONE;
                continue;
            }

            case SLOT_TASK_TYPE_RERANK: {
                const float* data = get_embedding_ptr(slot);
                float score = data ? data[0] : -1e6f;

                LOG_INFO("Rerank data: 0: %f", score);

                if (slot.rerank_current_index < slot.rerank_scores.size()) {
                    slot.rerank_scores[slot.rerank_current_index] = score;
                }

                slot.rerank_current_index++;

                if (slot.rerank_current_index < slot.rerank_prompt_tokens.size()) {
                    if (parent_ctx && parent_ctx->ctx) {
                        // Only this slot's sequence - a global clear would
                        // corrupt other slots' in-flight sequences
                        llama_memory_seq_rm(llama_get_memory(parent_ctx->ctx), slot.id, 0, -1);
                    }
                    slot.load_prompt(slot.rerank_prompt_tokens[slot.rerank_current_index]);
                    slot.state = SLOT_STATE_PROCESSING_PROMPT;
                    slot.i_batch = -1;
                    continue;
                }

                if (slot.on_rerank_callback) {
                    slot.on_rerank_callback(slot.request_id, slot.rerank_scores);
                }

                if (parent_ctx && parent_ctx->ctx) {
                    llama_memory_seq_rm(llama_get_memory(parent_ctx->ctx), slot.id, 0, -1);
                }

                slot.state = SLOT_STATE_DONE;
                continue;
            }

            default:
                LOG_ERROR("Slot %d: Unknown task type %d in sampling", slot.id, slot.task_type);
                slot.state = SLOT_STATE_DONE;
                continue;
        }
    }
}

// Release completed slots
void llama_rn_slot_manager::release_completed_slots() {
    for (auto& slot : slots) {
        if (slot.state == SLOT_STATE_DONE) {
            // Remove from active requests
            auto it = active_requests.find(slot.request_id);
            if (it != active_requests.end()) {
                active_requests.erase(it);
            }

            // Release slot
            release_slot(&slot);
        }
    }
}

// Main processing loop
void llama_rn_slot_manager::update_slots() {
    // Step 1: Terminalize cancellations and process pending queue (with mutex)
    bool completed_interrupted = false;
    {
        std::lock_guard<std::mutex> lock(slots_mutex);
        for (auto& slot : slots) {
            if (slot.is_interrupted &&
                (slot.state == SLOT_STATE_PROCESSING_PROMPT ||
                 slot.state == SLOT_STATE_GENERATING)) {
                LOG_INFO("Slot %d: Request interrupted", slot.id);
                complete_slot(slot);
                completed_interrupted = true;
            }
        }

        if (completed_interrupted) {
            release_completed_slots();
        }

        // Releasing interrupted slots first lets the next queued request claim
        // the slot in this update rather than waiting for another worker turn.
        process_pending_queue();
    }

    // Step 2: Check if any slots are active (with mutex)
    bool has_active = false;
    {
        std::lock_guard<std::mutex> lock(slots_mutex);
        for (const auto& slot : slots) {
            if (slot.state == SLOT_STATE_PROCESSING_PROMPT || slot.state == SLOT_STATE_GENERATING) {
                has_active = true;
                break;
            }
        }
    }

    if (!has_active) {
        // No active slots, return early after publishing worker-owned
        // cancellation cleanup to status subscribers.
        if (completed_interrupted) {
            bool has_subscribers = false;
            {
                std::lock_guard<std::mutex> lock(subscribers_mutex);
                has_subscribers = !status_subscribers.empty();
            }
            if (has_subscribers) {
                notify_status_change();
            }
        }
        return;
    }

    // Step 3: Build batch from all active slots (with mutex)
    {
        std::lock_guard<std::mutex> lock(slots_mutex);
        build_batch();
    }

    // Step 4: Process batch if we have tokens (NO mutex - llama_decode is thread-safe)
    if (batch.n_tokens > 0) {
        bool success = process_batch();
        if (!success) {
            LOG_ERROR("Batch processing failed");
            // Mark all active slots as done with error (with mutex)
            std::lock_guard<std::mutex> lock(slots_mutex);
            for (auto& slot : slots) {
                if (slot.state == SLOT_STATE_PROCESSING_PROMPT || slot.state == SLOT_STATE_GENERATING) {
                    slot.incomplete = true;
                    complete_slot(slot);
                }
            }
            release_completed_slots();
            return;
        }

        // Step 4.5: Calculate timing for slots that just finished prompt processing
        // This must happen AFTER batch has been decoded
        {
            std::lock_guard<std::mutex> lock(slots_mutex);
            const int64_t t_now = lm_ggml_time_us();
            for (auto& slot : slots) {
                if (slot.prompt_processing_finished) {
                    slot.t_start_generation = t_now;
                    slot.t_prompt_processing = (slot.t_start_generation - slot.t_start_process) / 1e6;
                    slot.prompt_processing_finished = false;  // Clear the flag

                    LOG_VERBOSE("Slot %d: Prompt processing complete, time=%.3fs, tokens=%d (cached=%d)",
                               slot.id, slot.t_prompt_processing, slot.n_prompt_tokens_processed, slot.n_prompt_tokens_cache);
                }
            }
        }
    }

    // Step 4.6: Save recurrent prompt checkpoints after decode (or if no decode needed)
    {
        std::lock_guard<std::mutex> lock(slots_mutex);
        for (auto& slot : slots) {
            if (!slot.save_prompt_state_pending || slot.save_prompt_state_tokens < 0) {
                continue;
            }

            if (slot.n_decoded > 0) {
                LOG_WARNING("Slot %d: Prompt checkpoint missed (generated tokens=%d), skipping",
                           slot.id, slot.n_decoded);
                slot.save_prompt_state_pending = false;
                continue;
            }

            if (slot.n_past >= slot.save_prompt_state_tokens) {
                const bool saved = slot.save_prompt_state_checkpoint();
                slot.save_prompt_state_pending = false;
                if (!saved) {
                    LOG_WARNING("Slot %d: Failed to save prompt checkpoint", slot.id);
                }
            }
        }
    }

    // Step 5: Terminalize cancellations, then sample and invoke callbacks for
    // remaining GENERATING slots (with mutex).
    {
        std::lock_guard<std::mutex> lock(slots_mutex);
        for (auto& slot : slots) {
            if (slot.is_interrupted &&
                (slot.state == SLOT_STATE_PROCESSING_PROMPT ||
                 slot.state == SLOT_STATE_GENERATING)) {
                LOG_INFO("Slot %d: Request interrupted", slot.id);
                complete_slot(slot);
            }
        }
        sample_and_callback();
    }

    // Step 6: Release completed slots (with mutex)
    {
        std::lock_guard<std::mutex> lock(slots_mutex);
        release_completed_slots();
    }

    // Step 7: Process pending queue again - assign requests to newly freed slots (with mutex)
    {
        std::lock_guard<std::mutex> lock(slots_mutex);
        process_pending_queue();
    }

    // Step 8: Notify subscribers of status change (outside of slots_mutex)
    // Check if there are subscribers before calling notify
    bool has_subscribers = false;
    {
        std::lock_guard<std::mutex> lock(subscribers_mutex);
        has_subscribers = !status_subscribers.empty();
    }
    if (has_subscribers) {
        notify_status_change();
    }
}

// Start background processing loop
void llama_rn_slot_manager::start_processing_loop() {
    // Check if already running
    if (processing_active.load()) {
        LOG_WARNING("Processing loop already active");
        return;
    }

    processing_active.store(true);

    // Start processing thread
    processing_thread = std::thread([this]() {
        LOG_INFO("Processing loop started");

        while (processing_active.load()) {
            // Call update_slots (protected by mutex)
            update_slots();

            // Wait for new work instead of sleeping
            // This efficiently blocks until notified or until there's work to do
            std::unique_lock<std::mutex> lock(slots_mutex);

            // Check if we have any active work or pending requests
            bool has_work = !queue_requests.empty();
            if (!has_work) {
                for (const auto& slot : slots) {
                    if (slot.state == SLOT_STATE_PROCESSING_PROMPT || slot.state == SLOT_STATE_GENERATING) {
                        has_work = true;
                        break;
                    }
                }
            }

            // If no work, wait for notification
            if (!has_work && processing_active.load()) {
                slots_cv.wait(lock, [this]() {
                    // Wake up if: there are pending requests, or processing should stop
                    return !queue_requests.empty() || !processing_active.load();
                });
            } else if (has_work) {
                lock.unlock();
                std::this_thread::sleep_for(std::chrono::microseconds(100));
            }
        }

        LOG_INFO("Processing loop stopped");
    });
}

// Stop background processing loop
void llama_rn_slot_manager::stop_processing_loop() {
    if (!processing_active.load()) {
        return;
    }

    LOG_INFO("Stopping processing loop...");
    processing_active.store(false);

    // Notify condition variable to wake up the thread
    slots_cv.notify_all();

    // Wait for processing thread to finish
    if (processing_thread.joinable()) {
        processing_thread.join();
    }

    LOG_INFO("Processing loop stopped");
}

// Get current parallel status
llama_rn_parallel_status llama_rn_slot_manager::get_status() {
    std::lock_guard<std::mutex> lock(slots_mutex);

    llama_rn_parallel_status status;
    status.n_parallel = n_parallel;
    status.active_slots = 0;
    status.queued_requests = static_cast<int32_t>(queue_requests.size());

    // Add active slot requests
    for (const auto& slot : slots) {
        if (slot.state != SLOT_STATE_IDLE && slot.state != SLOT_STATE_DONE) {
            status.active_slots++;

            llama_rn_request_status req_status;
            req_status.request_id = slot.request_id;

            // Map task type to string
            switch (slot.task_type) {
                case SLOT_TASK_TYPE_COMPLETION: req_status.type = "completion"; break;
                case SLOT_TASK_TYPE_EMBEDDING: req_status.type = "embedding"; break;
                case SLOT_TASK_TYPE_RERANK: req_status.type = "rerank"; break;
            }

            // Map state to string
            switch (slot.state) {
                case SLOT_STATE_IDLE: req_status.state = "idle"; break;
                case SLOT_STATE_PROCESSING_PROMPT: req_status.state = "processing_prompt"; break;
                case SLOT_STATE_GENERATING: req_status.state = "generating"; break;
                case SLOT_STATE_DONE: req_status.state = "done"; break;
            }

            req_status.prompt_length = slot.num_prompt_tokens;
            req_status.tokens_generated = slot.n_decoded;
            req_status.prompt_ms = slot.t_prompt_processing * 1e3;
            req_status.generation_ms = slot.t_token_generation * 1e3;
            req_status.tokens_per_second = (slot.n_decoded > 0 && slot.t_token_generation > 0.0)
                ? slot.n_decoded / slot.t_token_generation : 0.0;

            status.requests.push_back(req_status);
        }
    }

    // Add queued requests
    for (const auto& queued : queue_requests) {
        llama_rn_request_status req_status;
        req_status.request_id = queued.request_id;

        switch (queued.task_type) {
            case SLOT_TASK_TYPE_COMPLETION: req_status.type = "completion"; break;
            case SLOT_TASK_TYPE_EMBEDDING: req_status.type = "embedding"; break;
            case SLOT_TASK_TYPE_RERANK: req_status.type = "rerank"; break;
        }

        req_status.state = "queued";
        req_status.prompt_length = queued.prompt_tokens.size();
        req_status.tokens_generated = 0;
        req_status.prompt_ms = 0.0;
        req_status.generation_ms = 0.0;
        req_status.tokens_per_second = 0.0;

        status.requests.push_back(req_status);
    }

    return status;
}

bool llama_rn_slot_manager::has_pending_work() {
    std::lock_guard<std::mutex> lock(slots_mutex);

    if (!queue_requests.empty()) {
        return true;
    }

    for (const auto& slot : slots) {
        if (slot.state != SLOT_STATE_IDLE && slot.state != SLOT_STATE_DONE) {
            return true;
        }
    }

    return false;
}

// Notify all subscribers of status change
void llama_rn_slot_manager::notify_status_change() {
    // Get status snapshot first (acquires slots_mutex inside)
    llama_rn_parallel_status status = get_status();

    std::lock_guard<std::mutex> lock(subscribers_mutex);
    for (const auto& [id, callback] : status_subscribers) {
        if (callback) {
            callback(status);
        }
    }
}

// Add status subscriber
int32_t llama_rn_slot_manager::add_status_subscriber(
    std::function<void(const llama_rn_parallel_status&)> callback
) {
    std::lock_guard<std::mutex> lock(subscribers_mutex);
    int32_t id = next_subscriber_id++;
    status_subscribers[id] = callback;
    return id;
}

// Remove status subscriber
void llama_rn_slot_manager::remove_status_subscriber(int32_t subscriber_id) {
    std::lock_guard<std::mutex> lock(subscribers_mutex);
    status_subscribers.erase(subscriber_id);
}

} // namespace rnllama