#include "lemon/router.h"
#include "lemon/cloud_provider_registry.h"
#include "lemon/backends/cloud_server.h"
#include "lemon/backends/llamacpp_server.h"
#include "lemon/backends/fastflowlm_server.h"
#include "lemon/backends/ryzenaiserver.h"
#include "lemon/backends/whisper_server.h"
#include "lemon/backends/moonshine_server.h"
#include "lemon/backends/kokoro_server.h"
#include "lemon/backends/sd_server.h"
#include "lemon/backends/vllm_server.h"
#include "lemon/server_capabilities.h"
#include "lemon/error_types.h"
#include "lemon/recipe_options.h"
#include "lemon/auto_tune.h"
#include <iostream>
#include <algorithm>
#include "lemon/utils/aixlog.hpp"
#include "lemon/global_vram_monitor.h"
#include "lemon/eviction_engine.h"

namespace lemon {

Router::Router(RuntimeConfig* config, ModelManager* model_manager, BackendManager* backend_manager)
    : config_(config), model_manager_(model_manager), backend_manager_(backend_manager) {

    int max = config_->max_loaded_models();
    if (max == -1) {
    LOG(DEBUG, "Router") << "Max loaded models per type: unlimited" << std::endl;
    } else {
    LOG(DEBUG, "Router") << "Max loaded models per type: " << max << std::endl;
    }

    vram_monitor_ = std::make_unique<GlobalVramMonitor>();
    eviction_engine_ = std::make_unique<EvictionEngine>(this, vram_monitor_.get());

    // Always start the monitor/engine threads; they are cheap no-ops until the
    // user opts in. The monitor skips the VRAM poll when auto_evict is disabled,
    // and the engine's per-server check skips models that haven't opted in.
    // (auto_evict can be toggled at runtime via /internal/set, so we cannot gate
    // thread creation on the construction-time config value.)
    vram_monitor_->start();
    eviction_engine_->start();
}

Router::~Router() {
    LOG(DEBUG, "Router") << "Destructor: stopping monitors and unloading all models" << std::endl;
    if (eviction_engine_) eviction_engine_->stop();
    if (vram_monitor_) vram_monitor_->stop();
    unload_model("");  // Unload all
}

void Router::set_cloud_registry(CloudProviderRegistry* registry) {
    cloud_registry_ = registry;
}

WrappedServer* Router::find_server_by_model_name(const std::string& model_name) const {
    WrappedServer* unavailable_match = nullptr;
    for (const auto& server : loaded_servers_) {
        if (server->get_model_name() != model_name) {
            continue;
        }
        if (server->is_backend_alive()) {
            return server.get();
        }
        if (!unavailable_match) {
            unavailable_match = server.get();
        }
    }
    return unavailable_match;
}

std::string Router::resolve_model_name(const std::string& model_name) const {
    return model_name.empty() ? model_name : model_manager_->resolve_model_name(model_name);
}

WrappedServer* Router::get_most_recent_server() const {
    WrappedServer* most_recent = nullptr;
    for (const auto& server : loaded_servers_) {
        if (!server->is_backend_alive()) {
            continue;
        }
        if (!most_recent || server->get_last_access_time() > most_recent->get_last_access_time()) {
            most_recent = server.get();
        }
    }
    return most_recent;
}

void Router::prune_unavailable_servers_locked() {
    std::vector<WrappedServer*> unavailable;
    for (const auto& server : loaded_servers_) {
        if (!server->is_backend_alive()) {
            unavailable.push_back(server.get());
        }
    }

    for (auto* server : unavailable) {
        LOG(WARNING, "Router") << "Pruning unavailable backend for model: "
                                << server->get_model_name()
                                << " (state=" << server->get_backend_health_state() << ")"
                                << std::endl;
        evict_server(server);
    }
}

bool Router::is_watchdog_reset_response(const json& response) const {
    if (!response.is_object() || !response.contains("error") || !response["error"].is_object()) {
        return false;
    }

    const auto& error = response["error"];
    if (error.contains("code") && error["code"] == "backend_watchdog_reset") {
        return true;
    }
    if (error.contains("details") && error["details"].is_object() &&
        error["details"].contains("code") &&
        error["details"]["code"] == "backend_watchdog_reset") {
        return true;
    }
    return false;
}

bool Router::reload_model_after_watchdog_reset(const std::string& requested_model, const RecipeOptions& options) {
    try {
        LOG(WARNING, "Router") << "Reloading model after backend watchdog reset: "
                                << requested_model << std::endl;
        auto info = model_manager_->get_model_info(requested_model);
        bool was_pinned = false;
        {
            std::lock_guard<std::mutex> lock(load_mutex_);
            auto* existing = find_server_by_model_name(requested_model);
            if (existing) {
                was_pinned = existing->is_pinned();
            }
        }
        load_model(requested_model, info, options, true, false, was_pinned);
        return true;
    } catch (const std::exception& e) {
        LOG(ERROR, "Router") << "Automatic reload after watchdog reset failed for "
                              << requested_model << ": " << e.what() << std::endl;
        return false;
    }
}

int Router::count_servers_by_type(ModelType type) const {
    int count = 0;
    for (const auto& server : loaded_servers_) {
        // Cloud servers consume no local memory and stay loaded for free, so
        // they are excluded from the slot accounting that drives LRU eviction.
        if (server->get_recipe_options().get_recipe() == "cloud") {
            continue;
        }
        if (server->is_backend_alive() && server->get_model_type() == type) {
            count++;
        }
    }
    return count;
}

WrappedServer* Router::find_lru_server_by_type(ModelType type) const {
    WrappedServer* lru = nullptr;

    for (const auto& server : loaded_servers_) {
        // Cloud servers are not eviction candidates; they have no memory cost
        // and reloading them is essentially free, but evicting them throws
        // away the cached api key/upstream-id binding for no benefit.
        if (server->get_recipe_options().get_recipe() == "cloud") {
            continue;
        }
        if (server->is_backend_alive() && server->get_model_type() == type) {
            if (server->is_pinned()) {
                continue;
            }
            if (!lru || server->get_last_access_time() < lru->get_last_access_time()) {
                lru = server.get();
            }
        }
    }

    return lru;
}

bool Router::has_npu_server() const {
    for (const auto& server : loaded_servers_) {
        if (server->is_backend_alive() && (server->get_device_type() & DEVICE_NPU)) {
            return true;
        }
    }
    return false;
}

WrappedServer* Router::find_npu_server() const {
    for (const auto& server : loaded_servers_) {
        if (server->is_backend_alive() && (server->get_device_type() & DEVICE_NPU)) {
            return server.get();
        }
    }
    return nullptr;
}

WrappedServer* Router::find_npu_server_by_recipe(const std::string& recipe) const {
    for (const auto& server : loaded_servers_) {
        if (server->is_backend_alive() &&
            (server->get_device_type() & DEVICE_NPU) &&
            server->get_recipe_options().get_recipe() == recipe) {
            return server.get();
        }
    }
    return nullptr;
}

WrappedServer* Router::find_flm_server_by_type(ModelType type) const {
    for (const auto& server : loaded_servers_) {
        if (server->is_backend_alive() &&
            server->get_recipe_options().get_recipe() == "flm" &&
            server->get_model_type() == type) {
            return server.get();
        }
    }
    return nullptr;
}

// Helper: Evict all NPU servers
void Router::evict_all_npu_servers() {
    std::vector<WrappedServer*> npu_servers;
    for (const auto& server : loaded_servers_) {
        if (server->is_backend_alive() && (server->get_device_type() & DEVICE_NPU)) {
            npu_servers.push_back(server.get());
        }
    }
    for (auto* server : npu_servers) {
        LOG(INFO, "Router") << "Evicting NPU server: " << server->get_model_name() << std::endl;
        evict_server(server);
    }
}

void Router::evict_server(WrappedServer* server, int timeout_seconds) {
    if (!server) return;

    std::string model_name = server->get_model_name();
    LOG(INFO, "Router") << "Evicting model: " << model_name << std::endl;

    // Wait for any ongoing inference to complete. For watchdog-reset/dead
    // backends the wait is bounded so recovery can continue, but the object must
    // not be destroyed while a request thread may still be using this raw
    // pointer. In that case we leave a tombstoned server in loaded_servers_;
    // future prune/load calls will remove it once the request unwinds.
    const int wait_timeout = server->is_backend_alive() ? timeout_seconds : EVICTION_TIMEOUT;
    const bool idle = server->wait_until_not_busy(wait_timeout);
    if (!idle) {
        LOG(WARNING, "Router") << "Deferring eviction for model " << model_name
                                << " because requests are still unwinding after "
                                << EVICTION_TIMEOUT << "s (state="
                                << server->get_backend_health_state() << ")"
                                << std::endl;
        return;
    }

    server->unload();

    // Remove from vector
    loaded_servers_.erase(
        std::remove_if(loaded_servers_.begin(), loaded_servers_.end(),
                      [server](const std::unique_ptr<WrappedServer>& s) {
                          return s.get() == server;
                      }),
        loaded_servers_.end()
    );

    LOG(INFO, "Router") << "Evicted model: " << model_name << std::endl;
}

void Router::evict_all_servers() {
    LOG(INFO, "Router") << "Evicting all models (" << loaded_servers_.size() << " total)" << std::endl;

    // Copy raw pointers first; evict_server may erase entries and move
    // unique_ptrs inside the vector, but the pointed-to WrappedServer objects
    // remain stable until their individual eviction completes. Busy/dead
    // servers are safely left as tombstones for a later prune pass.
    std::vector<WrappedServer*> servers;
    servers.reserve(loaded_servers_.size());
    for (const auto& server : loaded_servers_) {
        servers.push_back(server.get());
    }

    for (auto* server : servers) {
        evict_server(server, EVICTION_TIMEOUT);
    }

    LOG(INFO, "Router") << "Evict all completed. Remaining tombstoned models: "
                         << loaded_servers_.size() << std::endl;
}

void Router::simulate_vram_pressure(double pct) {
    if (vram_monitor_) {
        vram_monitor_->simulate_pressure(pct);
    }
}

std::unique_ptr<WrappedServer> Router::create_backend_server(const ModelInfo& model_info) {
    std::unique_ptr<WrappedServer> new_server;
    std::string log_level = config_->log_level();

    if (model_info.recipe == "cloud") {
        LOG(DEBUG, "Router") << "Creating CloudServer backend (provider: "
                             << model_info.cloud_provider << ")" << std::endl;
        new_server = std::make_unique<backends::CloudServer>(model_info.cloud_provider, log_level,
                                                              model_manager_, backend_manager_,
                                                              cloud_registry_);
    } else if (model_info.recipe == "whispercpp") {
        LOG(DEBUG, "Router") << "Creating WhisperServer backend" << std::endl;
        new_server = std::make_unique<backends::WhisperServer>(log_level, model_manager_, backend_manager_);
    } else if (model_info.recipe == "moonshine") {
        LOG(DEBUG, "Router") << "Creating MoonshineServer backend" << std::endl;
        new_server = std::make_unique<backends::MoonshineServer>(log_level, model_manager_, backend_manager_);
    } else if (model_info.recipe == "kokoro") {
        LOG(DEBUG, "Router") << "Creating Kokoro backend" << std::endl;
        new_server = std::make_unique<backends::KokoroServer>(log_level, model_manager_, backend_manager_);
    } else if (model_info.recipe == "sd-cpp") {
        LOG(DEBUG, "Router") << "Creating SDServer backend" << std::endl;
        new_server = std::make_unique<backends::SDServer>(log_level, model_manager_, backend_manager_);
    } else if (model_info.recipe == "flm") {
        LOG(DEBUG, "Router") << "Creating FastFlowLM backend" << std::endl;
        new_server = std::make_unique<backends::FastFlowLMServer>(log_level, model_manager_, backend_manager_);
    } else if (model_info.recipe == "ryzenai-llm") {
        LOG(DEBUG, "Router") << "Creating RyzenAI-Server backend" << std::endl;

        std::string model_path = model_info.resolved_path();
        LOG(DEBUG, "Router") << "Using model path: " << model_path << std::endl;

        auto* ryzenai_server = new RyzenAIServer(model_info.model_name,
                                                  log_level == "debug", model_manager_, backend_manager_);
        ryzenai_server->set_model_path(model_path);
        new_server.reset(ryzenai_server);
    } else if (model_info.recipe == "vllm") {
        LOG(DEBUG, "Router") << "Creating vLLM backend" << std::endl;
        new_server = std::make_unique<backends::VLLMServer>(log_level, model_manager_, backend_manager_);
    } else {
        LOG(DEBUG, "Router") << "Creating LlamaCpp backend" << std::endl;
        new_server = std::make_unique<backends::LlamaCppServer>(log_level, model_manager_, backend_manager_);
    }

    return new_server;
}

void Router::load_model(const std::string& model_name,
                       const ModelInfo& model_info,
                       RecipeOptions options,
                       bool do_not_upgrade,
                       bool allow_reload_on_option_change,
                       std::optional<bool> pinned) {
    const std::string canonical_model_name = resolve_model_name(model_name);
    const std::string backend_option = model_info.recipe + "_backend";

    RecipeOptions tentative = options.inherit(model_info.recipe_options.inherit(
    RecipeOptions(model_info.recipe, config_->recipe_options(""))));
    json backend_json = tentative.get_option(backend_option);
    const std::string backend = backend_json.is_string() ? backend_json.get<std::string>() : "";

    // Second pass: rebuild defaults using the resolved backend
    RecipeOptions default_opt = RecipeOptions(model_info.recipe, config_->recipe_options(backend));
    RecipeOptions effective_options = options.inherit(model_info.recipe_options.inherit(default_opt));

    // LOAD SERIALIZATION STRATEGY (from spec: point #2 in Additional Considerations)
    std::unique_lock<std::mutex> lock(load_mutex_);

    // Wait if another thread is currently loading
    while (is_loading_) {
    LOG(INFO, "Router") << "Another load is in progress, waiting..." << std::endl;
        load_cv_.wait(lock);
    }

    // Mark that we're now loading (prevents concurrent loads)
    is_loading_ = true;

    LOG(DEBUG, "Router") << "Loading model: " << canonical_model_name
            << " (checkpoint: " << model_info.checkpoint()
            << ", recipe: " << model_info.recipe
            << ", type: " << model_type_to_string(model_info.type)
            << ", device: " << device_type_to_string(model_info.device) << ")" << std::endl;

    try {
        WrappedServer* existing_pre = find_server_by_model_name(canonical_model_name);
        bool final_pinned = false;
        if (pinned.has_value()) {
            final_pinned = pinned.value();
        } else if (existing_pre) {
            final_pinned = existing_pre->is_pinned();
        } else {
            final_pinned = (effective_options.get_option("pinned").is_boolean() && effective_options.get_option("pinned").get<bool>());
        }

        prune_unavailable_servers_locked();

        // Check if model is already loaded. Watchdog-reset or otherwise dead
        // entries are evicted first so auto-load performs a real lazy restart.
        WrappedServer* existing = find_server_by_model_name(canonical_model_name);
        if (existing && !existing->is_backend_alive()) {
            LOG(WARNING, "Router") << "Existing backend for " << canonical_model_name
                                    << " is unavailable (state="
                                    << existing->get_backend_health_state()
                                    << "), evicting before reload" << std::endl;
            evict_server(existing);
            existing = nullptr;
        }
        if (existing) {
            if (allow_reload_on_option_change &&
                existing->get_recipe_options().to_json() != effective_options.to_json()) {
                LOG(INFO, "Router") << "Options changed, reloading model: " << canonical_model_name << std::endl;
                evict_server(existing);
                // Fall through to create and load with new options
            } else {
                LOG(INFO, "Router") << "Model already loaded, updating access time and pinned status" << std::endl;
                existing->set_pinned(final_pinned);
                existing->update_access_time();
                is_loading_ = false;
                load_cv_.notify_all();
                return;
            }
        }

        // Determine model type and device
        ModelType model_type = model_info.type;
        DeviceType device_type = model_info.device;

        // Get max models for this type (same limit for all types)
        int max_models = config_->max_loaded_models();

        // NPU EXCLUSIVITY CHECK (recipe-aware rules)
        // FLM can run up to 3 concurrent NPU processes (1 LLM + 1 transcription + 1 embedding)
        // RyzenAI and WhisperCpp lock the entire NPU exclusively
        if (device_type & DEVICE_NPU) {
            if (model_info.recipe == "ryzenai-llm" || model_info.recipe == "whispercpp") {
                // Exclusive NPU recipes - evict ALL NPU servers
                if (has_npu_server()) {
                    LOG(INFO, "Router") << model_info.recipe
                              << " requires exclusive NPU access, evicting all NPU servers..." << std::endl;
                    evict_all_npu_servers();
                }
            } else if (model_info.recipe == "flm") {
                // FLM can coexist with other FLM types, but not with exclusive-NPU recipes
                // 1. Evict any exclusive-NPU server (mutually exclusive)
                for (const std::string& exclusive_recipe : {"ryzenai-llm", "whispercpp"}) {
                    WrappedServer* exclusive_server = find_npu_server_by_recipe(exclusive_recipe);
                    if (exclusive_server) {
                        LOG(INFO, "Router") << "FLM cannot coexist with " << exclusive_recipe
                                  << ", evicting: " << exclusive_server->get_model_name() << std::endl;
                        evict_server(exclusive_server);
                    }
                }
                // 2. Evict FLM of the SAME model type (max 1 per type: 1 LLM, 1 transcription, 1 embed)
                WrappedServer* same_type_flm = find_flm_server_by_type(model_type);
                if (same_type_flm) {
                    LOG(INFO, "Router") << "FLM " << model_type_to_string(model_type)
                              << " slot occupied by: " << same_type_flm->get_model_name()
                              << ", evicting..." << std::endl;
                    evict_server(same_type_flm);
                }
            } else {
                // Unknown NPU recipe - default to exclusive access
                if (has_npu_server()) {
                    LOG(INFO, "Router") << "Unknown NPU recipe, evicting all NPU servers..." << std::endl;
                    evict_all_npu_servers();
                }
            }
        }

        // LRU EVICTION CHECK (from spec: Least Recently Used Cache)
        // Skip eviction if unlimited (-1). Cloud-recipe loads also skip the
        // check entirely: they consume no local resources, so they have no
        // business kicking a warm local model out of memory.
        bool is_cloud_load = (model_info.recipe == "cloud");
        int current_count = count_servers_by_type(model_type);
        if (!is_cloud_load && max_models != -1 && current_count >= max_models) {
            WrappedServer* lru = find_lru_server_by_type(model_type);
            if (lru) {
                LOG(INFO, "Router") << "Slot limit reached for type "
                          << model_type_to_string(model_type)
                          << ", evicting LRU: " << lru->get_model_name() << std::endl;
                evict_server(lru);
            } else {
                is_loading_ = false;
                load_cv_.notify_all();
                throw SlotsPinnedException(model_type_to_string(model_type));
            }
        }

        // Auto-tune: resolve ctx_size = -1 → computed from memory + arch metadata
        // Done AFTER eviction so that freed VRAM/RAM is visible to the memory query.
        int64_t auto_ctx = resolve_auto_ctx_size(effective_options, model_info);
        if (auto_ctx > 0) {
            LOG(INFO, "Router") << "Auto-tune ctx_size resolved to " << auto_ctx << std::endl;
            effective_options.set_option("ctx_size", auto_ctx);
        }

        LOG(DEBUG, "Router") << "Effective settings: " << effective_options.to_log_string() << std::endl;

        // Create new backend server
        std::unique_ptr<WrappedServer> new_server = create_backend_server(model_info);

        // Set model metadata
        new_server->set_model_metadata(canonical_model_name, model_info.checkpoint(), model_type, device_type, effective_options);
        new_server->set_pinned(final_pinned);
        new_server->update_access_time();

        // CRITICAL: Release lock before slow backend startup
        lock.unlock();

        // Load the backend (this can take 30-60 seconds)
        LOG(DEBUG, "Router") << "Starting backend (this may take a moment)..." << std::endl;
        bool load_success = false;
        std::string error_message;
        auto load_start = std::chrono::steady_clock::now();

        try {
            new_server->load(canonical_model_name, model_info, effective_options, do_not_upgrade);
            load_success = true;
            auto load_end = std::chrono::steady_clock::now();
            new_server->set_load_duration_ms(std::chrono::duration_cast<std::chrono::milliseconds>(load_end - load_start).count());
            LOG(DEBUG, "Router") << "Backend started successfully in " << new_server->get_load_duration_ms() << "ms" << std::endl;
        } catch (const std::exception& e) {
            error_message = e.what();
            load_success = false;
            LOG(ERROR, "Router") << "Backend load failed: " << error_message << std::endl;
        }

        lock.lock();

        if (load_success) {
            // Success: Refresh access time so this model is returned by
            // get_most_recent_server() (the pre-load timestamp from line 316
            // may have been overtaken by other models serving requests while
            // the lock was released during the slow backend load).
            new_server->update_access_time();
            new_server->set_state(ModelState::READY);

            // Add to loaded servers
            loaded_servers_.push_back(std::move(new_server));

            is_loading_ = false;
            load_cv_.notify_all();

            LOG(INFO, "Router") << "Model loaded successfully. Total loaded: "
                      << loaded_servers_.size() << std::endl;
        } else {
            // ERROR HANDLING (from spec: Error Handling section)
            // Check if error is "file not found" (exception to nuclear policy)
            bool is_file_not_found = (error_message.find("not found") != std::string::npos ||
                                     error_message.find("does not exist") != std::string::npos ||
                                     error_message.find("No such file") != std::string::npos);

            is_loading_ = false;
            load_cv_.notify_all();

            if (is_file_not_found) {
                LOG(ERROR, "Router") << "File not found error, NOT evicting other models" << std::endl;
                throw std::runtime_error(error_message);
            }

            // Nuclear option: evict all models and retry
            LOG(WARNING, "Router") << "Load failed with non-file-not-found error, "
                      << "evicting all models and retrying..." << std::endl;

            evict_all_servers();

            // Mark loading again for retry
            is_loading_ = true;

            // Create new server for retry
            std::unique_ptr<WrappedServer> retry_server = create_backend_server(model_info);
            retry_server->set_model_metadata(canonical_model_name, model_info.checkpoint(), model_type, device_type, effective_options);
            retry_server->set_pinned(final_pinned);
            retry_server->update_access_time();

            lock.unlock();

            LOG(DEBUG, "Router") << "Retrying backend load..." << std::endl;
            try {
                auto retry_start = std::chrono::steady_clock::now();
                retry_server->load(canonical_model_name, model_info, effective_options, do_not_upgrade);
                auto retry_end = std::chrono::steady_clock::now();
                retry_server->set_load_duration_ms(std::chrono::duration_cast<std::chrono::milliseconds>(retry_end - retry_start).count());

                lock.lock();

                retry_server->set_state(ModelState::READY);
                loaded_servers_.push_back(std::move(retry_server));
                is_loading_ = false;
                load_cv_.notify_all();

                LOG(DEBUG, "Router") << "Retry successful in " << retry_server->get_load_duration_ms() << "ms!" << std::endl;
            } catch (const std::exception& retry_error) {
                lock.lock();
                is_loading_ = false;
                load_cv_.notify_all();

                LOG(ERROR, "Router") << "Retry also failed: " << retry_error.what() << std::endl;
                throw;
            }
        }

    } catch (const std::exception& e) {
        LOG(ERROR, "Router") << "Failed to load model: " << e.what() << std::endl;

        if (!lock.owns_lock()) {
            lock.lock();
        }
        is_loading_ = false;
        load_cv_.notify_all();

        throw;
    }
}

void Router::unload_model(const std::string& model_name) {
    std::lock_guard<std::mutex> lock(load_mutex_);

    if (model_name.empty()) {
        // Unload all models
        LOG(INFO, "Router") << "Unload all models called" << std::endl;
        evict_all_servers();
    } else {
        // Unload specific model
        LOG(INFO, "Router") << "Unload model called: " << model_name << std::endl;
        std::string canonical_model_name = resolve_model_name(model_name);
        WrappedServer* server = find_server_by_model_name(canonical_model_name);
        if (!server) {
            throw std::runtime_error("Model not loaded: " + model_name);
        }
        evict_server(server);
    }
}

void Router::evict_if_committed(const std::string& model_name) {
    std::lock_guard<std::mutex> lock(load_mutex_);

    WrappedServer* server = find_server_by_model_name(model_name);
    if (!server) {
        return;  // Already gone
    }

    // Atomically confirm the model is still idle and EVICTING. If a request
    // rescued it (now IN_USE) this returns false and reverts it to READY, so we
    // leave it loaded — no crashed generation, no talking to a dead subprocess.
    if (!server->try_commit_eviction()) {
        LOG(INFO, "Router") << "Eviction of " << model_name
                            << " cancelled (rescued by in-flight request)" << std::endl;
        return;
    }

    evict_server(server);
}

std::string Router::get_loaded_model() const {
    std::lock_guard<std::mutex> lock(load_mutex_);
    WrappedServer* server = get_most_recent_server();
    return server ? model_manager_->get_public_model_name(server->get_model_name()) : "";
}

std::string Router::get_loaded_recipe() const {
    std::lock_guard<std::mutex> lock(load_mutex_);
    WrappedServer* server = get_most_recent_server();
    if (!server) return "";

    // Get the actual recipe from the server's recipe options
    return server->get_recipe_options().get_recipe();
}

json Router::get_all_loaded_models() const {
    std::lock_guard<std::mutex> lock(load_mutex_);

    json result = json::array();

    for (const auto& server : loaded_servers_) {
        const bool backend_alive = server->is_backend_alive();
        if (!backend_alive) {
            continue;
        }

        json model_info;
        model_info["model_name"] = model_manager_->get_public_model_name(server->get_model_name());
        model_info["checkpoint"] = server->get_checkpoint();
        model_info["type"] = model_type_to_string(server->get_model_type());
        model_info["device"] = device_type_to_string(server->get_device_type());
        model_info["backend_url"] = server->get_address();  // For debugging port issues
        model_info["pid"] = server->get_process_id();
        model_info["status"] = model_state_to_string(server->get_state());
        model_info["backend_alive"] = true;
        model_info["backend_health"] = server->get_backend_health_state();
        model_info["loaded"] = true;
        model_info["watchdog_reset"] = server->was_watchdog_triggered();
        std::string watchdog_reason = server->get_watchdog_reset_reason();
        if (!watchdog_reason.empty()) {
            model_info["watchdog_reset_reason"] = watchdog_reason;
        }
        model_info["pinned"] = server->is_pinned();
        RecipeOptions recipe_options =  server->get_recipe_options();
        model_info["recipe"] = recipe_options.get_recipe();
        model_info["recipe_options"] = recipe_options.to_json();

        // Static metadata from the registry entry. Cloud models carry the
        // provider-reported context window + per-million-token cost (recorded
        // at discovery by ModelManager::refresh_cloud_models); local models
        // surface their runtime context via recipe_options instead.
        try {
            const ModelInfo reg_info = model_manager_->get_model_info(server->get_model_name());
            if (reg_info.max_context_window > 0) {
                model_info["max_context_window"] = reg_info.max_context_window;
            }
            if (reg_info.cost_input_per_million >= 0) {
                model_info["cost_input_per_million"] = reg_info.cost_input_per_million;
            }
            if (reg_info.cost_output_per_million >= 0) {
                model_info["cost_output_per_million"] = reg_info.cost_output_per_million;
            }
        } catch (...) {
            // Registry entry not found (raced with a delete) — skip static metadata.
        }

        // Convert timestamp to milliseconds since epoch
        auto time_point = server->get_last_access_time();
        auto duration = time_point.time_since_epoch();
        auto millis = std::chrono::duration_cast<std::chrono::milliseconds>(duration).count();
        model_info["last_use"] = millis;

        result.push_back(model_info);
    }

    return result;
}

json Router::get_max_model_limits() const {
    int max = config_->max_loaded_models();
    return {
        {"llm", max},
        {"embedding", max},
        {"reranking", max},
        {"transcription", max},
        {"image", max},
        {"tts", max}
    };
}

bool Router::is_model_loaded() const {
    std::lock_guard<std::mutex> lock(load_mutex_);
    for (const auto& server : loaded_servers_) {
        if (server->is_backend_alive()) {
            return true;
        }
    }
    return false;
}

bool Router::is_model_loaded(const std::string& model_name) const {
    std::lock_guard<std::mutex> lock(load_mutex_);
    auto* server = find_server_by_model_name(resolve_model_name(model_name));
    return server != nullptr && server->is_backend_alive();
}

RecipeOptions Router::get_model_recipe_options(const std::string& model_name) const {
    std::lock_guard<std::mutex> lock(load_mutex_);
    auto* server = find_server_by_model_name(resolve_model_name(model_name));
    if (server && server->is_backend_alive()) return server->get_recipe_options();
    return RecipeOptions();
}

ModelType Router::get_model_type(const std::string& model_name) const {
    std::lock_guard<std::mutex> lock(load_mutex_);
    WrappedServer* server = model_name.empty()
        ? get_most_recent_server()
        : find_server_by_model_name(resolve_model_name(model_name));
    return (server && server->is_backend_alive()) ? server->get_model_type() : ModelType::LLM;
}

std::string Router::get_backend_address() const {
    std::lock_guard<std::mutex> lock(load_mutex_);
    WrappedServer* server = get_most_recent_server();
    return (server && server->is_backend_alive()) ? server->get_address() : "";
}

std::string Router::get_streaming_transcription_address(const std::string& model_name) const {
    std::lock_guard<std::mutex> lock(load_mutex_);
    WrappedServer* server = nullptr;
    if (!model_name.empty()) {
        // Route by the session's requested model, like the normal inference
        // path — most-recent would misroute multi-model setups (e.g. a
        // Whisper session connecting to Moonshine's stream)
        server = find_server_by_model_name(resolve_model_name(model_name));
    } else {
        server = get_most_recent_server();
    }
    if (!server) {
        return "";
    }
    auto* streaming = dynamic_cast<IStreamingTranscriptionServer*>(server);
    return streaming ? streaming->get_streaming_address() : "";
}

template<typename Func>
auto Router::execute_inference(const json& request, Func&& inference_func) -> decltype(inference_func(nullptr)) {
    std::string requested_model;
    if (request.contains("model") && request["model"].is_string()) {
        requested_model = request["model"].get<std::string>();
    }

    if (requested_model.empty()) {
        return ErrorResponse::from_exception(InvalidRequestException("No model specified in request"));
    }

    // A watchdog reset should be transparent for non-streaming calls when the
    // backend died before any response was returned. Retry exactly once after a
    // lazy reload; streaming paths deliberately do not retry after partial data.
    for (int attempt = 0; attempt < 2; ++attempt) {
        WrappedServer* server = nullptr;
        RecipeOptions restart_options;
        std::string restart_model_name;
        bool should_reload_before_request = false;

        {
            std::lock_guard<std::mutex> lock(load_mutex_);
            server = find_server_by_model_name(resolve_model_name(requested_model));
            if (!server) {
                return ErrorResponse::from_exception(ModelNotLoadedException(requested_model));
            }

            if (!server->is_backend_alive()) {
                restart_options = server->get_recipe_options();
                restart_model_name = server->get_model_name();
                should_reload_before_request = true;
            } else {
                if (!server->acquire_for_inference()) {
                    return ErrorResponse::from_exception(ModelNotLoadedException(requested_model));
                }
                server->update_access_time();
            }
        } // Lock released here

        if (should_reload_before_request) {
            if (restart_model_name.empty()) {
                restart_model_name = requested_model;
            }
            if (attempt == 0 && reload_model_after_watchdog_reset(restart_model_name, restart_options)) {
                continue;
            }
            return ErrorResponse::create(
                "Backend for model '" + requested_model + "' is unavailable",
                ErrorType::BACKEND_ERROR,
                {{"code", "backend_unavailable"}, {"retryable", true}}
            );
        }

        try {
            auto response = inference_func(server);
            const bool watchdog_reset =
                server->was_watchdog_triggered() || is_watchdog_reset_response(response);

            if (attempt == 0 && watchdog_reset) {
                restart_options = server->get_recipe_options();
                restart_model_name = server->get_model_name();
            }

            server->release_inference();

            if (attempt == 0 && watchdog_reset) {
                if (restart_model_name.empty()) {
                    restart_model_name = requested_model;
                }
                if (reload_model_after_watchdog_reset(restart_model_name, restart_options)) {
                    continue;
                }
            }

            return response;
        } catch (...) {
            server->release_inference();
            throw;
        }
    }

    return ErrorResponse::create(
        "Backend watchdog reset recovery failed for model '" + requested_model + "'",
        ErrorType::BACKEND_ERROR,
        {{"code", "backend_watchdog_reset"}, {"retryable", true}}
    );
}

// Template method for streaming execution
template<typename Func>
void Router::execute_streaming(const std::string& request_body, httplib::DataSink& sink, Func&& streaming_func) {
    WrappedServer* server = nullptr;
    std::string requested_model;

    try {
        json request = json::parse(request_body);
        if (request.contains("model") && request["model"].is_string()) {
            requested_model = request["model"].get<std::string>();
        }
    } catch (...) {
        LOG(DEBUG, "Router") << "Failed to parse request body for model extraction" << std::endl;
    }

    if (requested_model.empty()) {
        LOG(ERROR, "Router") << "No model specified in streaming request" << std::endl;
        json error = ErrorResponse::from_exception(InvalidRequestException("No model specified in request"));
        std::string error_msg = "data: " + error.dump() + "\n\n";
        sink.write(error_msg.c_str(), error_msg.size());
        sink.done();
        return;
    }

    for (int attempt = 0; attempt < 2; ++attempt) {
        RecipeOptions restart_options;
        std::string restart_model_name;
        bool should_reload_before_request = false;

        {
            std::lock_guard<std::mutex> lock(load_mutex_);
            server = find_server_by_model_name(resolve_model_name(requested_model));
            if (!server) {
                json error = ErrorResponse::from_exception(ModelNotLoadedException(requested_model));
                std::string error_msg = "data: " + error.dump() + "\n\n";
                sink.write(error_msg.c_str(), error_msg.size());
                sink.done();
                return;
            }

            if (!server->is_backend_alive()) {
                restart_options = server->get_recipe_options();
                restart_model_name = server->get_model_name();
                should_reload_before_request = true;
            } else {
                if (!server->acquire_for_inference()) {
                    std::string error_msg =
                        "data: {\"error\":{\"message\":\"Model evicted: " + requested_model +
                        "\",\"type\":\"model_not_loaded\"}}\n\n";
                    sink.write(error_msg.c_str(), error_msg.size());
                    sink.done();
                    return;
                }
                server->update_access_time();
            }
        } // Lock released here

        if (should_reload_before_request) {
            if (restart_model_name.empty()) {
                restart_model_name = requested_model;
            }
            if (attempt == 0 && reload_model_after_watchdog_reset(restart_model_name, restart_options)) {
                continue;
            }

            json error = ErrorResponse::create(
                "Backend for model '" + requested_model + "' is unavailable",
                ErrorType::BACKEND_ERROR,
                {{"code", "backend_unavailable"}, {"retryable", true}}
            );
            std::string error_msg = "data: " + error.dump() + "\n\n";
            sink.write(error_msg.c_str(), error_msg.size());
            sink.done();
            return;
        }

        try {
            streaming_func(server);
            const bool watchdog_reset = server->was_watchdog_triggered();

            if (watchdog_reset) {
                restart_options = server->get_recipe_options();
                restart_model_name = server->get_model_name();
            }

            server->release_inference();

            // Do not replay a streaming response after bytes may have reached the
            // client. Reload immediately so the next request does not see a
            // stale tombstone, then return the stream outcome as-is.
            if (watchdog_reset) {
                if (restart_model_name.empty()) {
                    restart_model_name = requested_model;
                }
                reload_model_after_watchdog_reset(restart_model_name, restart_options);
            }
            return;
        } catch (const BackendStreamRetryableReset& e) {
            restart_options = server->get_recipe_options();
            restart_model_name = server->get_model_name();
            server->release_inference();

            if (restart_model_name.empty()) {
                restart_model_name = requested_model;
            }
            if (attempt == 0 && reload_model_after_watchdog_reset(restart_model_name, restart_options)) {
                continue;
            }

            json error = ErrorResponse::create(
                std::string("Backend for model '") + requested_model +
                    "' crashed before streaming started and could not be reloaded: " + e.what(),
                ErrorType::BACKEND_ERROR,
                {{"code", "backend_watchdog_reset"}, {"retryable", true}}
            );
            std::string error_msg = "data: " + error.dump() + "\n\n";
            sink.write(error_msg.c_str(), error_msg.size());
            sink.done();
            return;
        } catch (...) {
            server->release_inference();
            throw;
        }
    }
}

json Router::chat_completion(const json& request) {
    return execute_inference(request, [&](WrappedServer* server) {
        return server->chat_completion(request);
    });
}

json Router::completion(const json& request) {
    return execute_inference(request, [&](WrappedServer* server) {
        return server->completion(request);
    });
}

json Router::embeddings(const json& request) {
    return execute_inference(request, [&](WrappedServer* server) {
        auto embeddings_server = dynamic_cast<IEmbeddingsServer*>(server);
        if (!embeddings_server) {
            return ErrorResponse::from_exception(
                UnsupportedOperationException("Embeddings", device_type_to_string(server->get_device_type()))
            );
        }
        return embeddings_server->embeddings(request);
    });
}

json Router::reranking(const json& request) {
    return execute_inference(request, [&](WrappedServer* server) {
        auto reranking_server = dynamic_cast<IRerankingServer*>(server);
        if (!reranking_server) {
            return ErrorResponse::from_exception(
                UnsupportedOperationException("Reranking", device_type_to_string(server->get_device_type()))
            );
        }
        return reranking_server->reranking(request);
    });
}

json Router::get_slots() {
    WrappedServer* server = nullptr;
    ISlotsServer* slots_server = nullptr;

    {
        std::lock_guard<std::mutex> lock(load_mutex_);
        server = get_most_recent_server();
        if (!server) {
            return ErrorResponse::from_exception(
                ModelNotLoadedException("No models loaded")
            );
        }

        // Check if server supports slots capability
        slots_server = dynamic_cast<ISlotsServer*>(server);
        if (!slots_server) {
            return ErrorResponse::from_exception(
                UnsupportedOperationException("Slots", device_type_to_string(server->get_device_type()))
            );
        }

        // Mark as busy and update access time
        if (!server->acquire_for_inference()) {
            return ErrorResponse::from_exception(ModelNotLoadedException("No models loaded"));
        }
        server->update_access_time();
    } // Lock released here

    // Execute without holding lock (but busy flag prevents eviction)
    try {
        auto response = slots_server->get_slots();
        server->release_inference();
        return response;
    } catch (...) {
        server->release_inference();
        throw;
    }
}

json Router::slots_action(int slot_id, const std::string& action, const json& request_body) {
    WrappedServer* server = nullptr;
    ISlotsServer* slots_server = nullptr;

    {
        std::lock_guard<std::mutex> lock(load_mutex_);
        server = get_most_recent_server();
        if (!server) {
            return ErrorResponse::from_exception(
                ModelNotLoadedException("No models loaded")
            );
        }

        // Check if server supports slots capability
        slots_server = dynamic_cast<ISlotsServer*>(server);
        if (!slots_server) {
            return ErrorResponse::from_exception(
                UnsupportedOperationException("Slots", device_type_to_string(server->get_device_type()))
            );
        }

        // Mark as busy and update access time
        if (!server->acquire_for_inference()) {
            return ErrorResponse::from_exception(ModelNotLoadedException("No models loaded"));
        }
        server->update_access_time();
    } // Lock released here

    // Execute without holding lock (but busy flag prevents eviction)
    try {
        auto response = slots_server->slots_action(slot_id, action, request_body);
        server->release_inference();
        return response;
    } catch (...) {
        server->release_inference();
        throw;
    }
}

json Router::tokenize(const json& request_body) {
    WrappedServer* server = nullptr;
    ITokenizerServer* tokenizer_server = nullptr;

    {
        std::lock_guard<std::mutex> lock(load_mutex_);
        server = get_most_recent_server();
        if (!server) {
            return ErrorResponse::from_exception(
                ModelNotLoadedException("No models loaded")
            );
        }

        // Check if server supports tokenize capability
        tokenizer_server = dynamic_cast<ITokenizerServer*>(server);
        if (!tokenizer_server) {
            return ErrorResponse::from_exception(
                UnsupportedOperationException("Tokenization", device_type_to_string(server->get_device_type()))
            );
        }

        // Mark as busy and update access time
        if (!server->acquire_for_inference()) {
            return ErrorResponse::from_exception(ModelNotLoadedException("No models loaded"));
        }
        server->update_access_time();
    } // Lock released here

    // Execute without holding lock (but busy flag prevents eviction)
    try {
        auto response = tokenizer_server->tokenize(request_body);
        server->release_inference();
        return response;
    } catch (...) {
        server->release_inference();
        throw;
    }
}

json Router::responses(const json& request) {
    return execute_inference(request, [&](WrappedServer* server) {
        return server->responses(request);
    });
}

json Router::audio_transcriptions(const json& request) {
    return execute_inference(request, [&](WrappedServer* server) {
        auto transcription_server = dynamic_cast<ITranscriptionServer*>(server);
        if (!transcription_server) {
            return ErrorResponse::from_exception(
                UnsupportedOperationException("Audio transcription", device_type_to_string(server->get_device_type()))
            );
        }
        return transcription_server->audio_transcriptions(request);
    });
}

void Router::audio_speech(const json& request, httplib::DataSink& sink) {
    execute_streaming(request.dump(), sink, [&](WrappedServer* server) {
        auto tts_server = dynamic_cast<ITextToSpeechServer*>(server);
        if (!tts_server) {
            throw UnsupportedOperationException("Text to speech", device_type_to_string(server->get_device_type()));
        }
        tts_server->audio_speech(request, sink);
    });
}

json Router::image_generations(const json& request) {
    return execute_inference(request, [&](WrappedServer* server) {
        auto image_server = dynamic_cast<IImageServer*>(server);
        if (!image_server) {
            return ErrorResponse::from_exception(
                UnsupportedOperationException("Image generation", device_type_to_string(server->get_device_type()))
            );
        }
        return image_server->image_generations(request);
    });
}

json Router::image_edits(const json& request) {
    return execute_inference(request, [&](WrappedServer* server) {
        auto image_server = dynamic_cast<IImageServer*>(server);
        if (!image_server) {
            return ErrorResponse::from_exception(
                UnsupportedOperationException("Image editing", device_type_to_string(server->get_device_type()))
            );
        }
        return image_server->image_edits(request);
    });
}

json Router::image_variations(const json& request) {
    return execute_inference(request, [&](WrappedServer* server) {
        auto image_server = dynamic_cast<IImageServer*>(server);
        if (!image_server) {
            return ErrorResponse::from_exception(
                UnsupportedOperationException("Image variations", device_type_to_string(server->get_device_type()))
            );
        }
        return image_server->image_variations(request);
    });
}

json Router::get_stats() const {
    std::lock_guard<std::mutex> lock(telemetry_mutex_);
    return aggregate_telemetry_.to_json();
}

json Router::get_metrics_snapshot() const {
    json result;
    result["loaded_models"] = json::array();
    result["model_metrics"] = json::array();
    result["totals"] = {
        {"requests", 0},
        {"input_tokens", 0},
        {"output_tokens", 0},
        {"prompt_tokens", 0}
    };

    std::map<std::string, ModelTelemetryIdentity> loaded_identities;

    {
        std::lock_guard<std::mutex> lock(load_mutex_);
        for (const auto& server : loaded_servers_) {
            ModelTelemetryIdentity identity = get_telemetry_identity(server.get());
            loaded_identities[identity.key()] = identity;

            json model_info;
            model_info["model_name"] = model_manager_->get_public_model_name(identity.model_name);
            model_info["checkpoint"] = identity.checkpoint;
            model_info["type"] = identity.type;
            model_info["device"] = identity.device;
            model_info["backend_url"] = server->get_address();
            model_info["pid"] = server->get_process_id();
            model_info["recipe"] = identity.recipe;
            result["loaded_models"].push_back(model_info);
        }
    }

    {
        std::lock_guard<std::mutex> lock(telemetry_mutex_);
        for (const auto& item : telemetry_by_model_) {
            const auto& record = item.second;
            json model_info;
            model_info["model_name"] = model_manager_->get_public_model_name(record.identity.model_name);
            model_info["checkpoint"] = record.identity.checkpoint;
            model_info["type"] = record.identity.type;
            model_info["device"] = record.identity.device;
            model_info["recipe"] = record.identity.recipe;
            model_info["loaded"] = loaded_identities.find(item.first) != loaded_identities.end();
            model_info["telemetry"] = record.telemetry.to_json();
            result["model_metrics"].push_back(model_info);
        }

        for (const auto& item : loaded_identities) {
            if (telemetry_by_model_.find(item.first) != telemetry_by_model_.end()) {
                continue;
            }
            const auto& identity = item.second;
            json model_info;
            model_info["model_name"] = model_manager_->get_public_model_name(identity.model_name);
            model_info["checkpoint"] = identity.checkpoint;
            model_info["type"] = identity.type;
            model_info["device"] = identity.device;
            model_info["recipe"] = identity.recipe;
            model_info["loaded"] = true;
            model_info["telemetry"] = Telemetry().to_json();
            result["model_metrics"].push_back(model_info);
        }

        result["totals"]["requests"] = aggregate_telemetry_.request_count_total;
        result["totals"]["input_tokens"] = aggregate_telemetry_.input_tokens_total;
        result["totals"]["output_tokens"] = aggregate_telemetry_.output_tokens_total;
        result["totals"]["prompt_tokens"] = aggregate_telemetry_.prompt_tokens_total;
    }

    return result;
}

ModelTelemetryIdentity Router::get_telemetry_identity(WrappedServer* server) const {
    if (!server) {
        return {};
    }

    RecipeOptions recipe_options = server->get_recipe_options();
    return {
        server->get_model_name(),
        server->get_checkpoint(),
        model_type_to_string(server->get_model_type()),
        device_type_to_string(server->get_device_type()),
        recipe_options.get_recipe()
    };
}

void Router::record_telemetry_for_model(const ModelTelemetryIdentity& identity,
                                        int input_tokens,
                                        int output_tokens,
                                        double time_to_first_token,
                                        double tokens_per_second) {
    if (identity.model_name.empty()) {
        return;
    }

    std::lock_guard<std::mutex> lock(telemetry_mutex_);
    ModelTelemetryRecord& record = telemetry_by_model_[identity.key()];
    record.identity = identity;
    Telemetry& model_telemetry = record.telemetry;
    model_telemetry.input_tokens = input_tokens;
    model_telemetry.output_tokens = output_tokens;
    model_telemetry.time_to_first_token = time_to_first_token;
    model_telemetry.tokens_per_second = tokens_per_second;
    model_telemetry.request_count_total++;

    aggregate_telemetry_.input_tokens = input_tokens;
    aggregate_telemetry_.output_tokens = output_tokens;
    aggregate_telemetry_.time_to_first_token = time_to_first_token;
    aggregate_telemetry_.tokens_per_second = tokens_per_second;
    aggregate_telemetry_.request_count_total++;

    if (input_tokens > 0) {
        model_telemetry.input_tokens_total += static_cast<uint64_t>(input_tokens);
        aggregate_telemetry_.input_tokens_total += static_cast<uint64_t>(input_tokens);
    }
    if (output_tokens > 0) {
        model_telemetry.output_tokens_total += static_cast<uint64_t>(output_tokens);
        aggregate_telemetry_.output_tokens_total += static_cast<uint64_t>(output_tokens);
    }
}

void Router::record_prompt_tokens_for_model(const ModelTelemetryIdentity& identity, int prompt_tokens) {
    if (identity.model_name.empty()) {
        return;
    }

    std::lock_guard<std::mutex> lock(telemetry_mutex_);
    ModelTelemetryRecord& record = telemetry_by_model_[identity.key()];
    record.identity = identity;
    Telemetry& model_telemetry = record.telemetry;
    model_telemetry.prompt_tokens = prompt_tokens;
    aggregate_telemetry_.prompt_tokens = prompt_tokens;
    if (prompt_tokens > 0) {
        model_telemetry.prompt_tokens_total += static_cast<uint64_t>(prompt_tokens);
        aggregate_telemetry_.prompt_tokens_total += static_cast<uint64_t>(prompt_tokens);
    }
}

void Router::update_telemetry(const std::string& model_name,
                              int input_tokens, int output_tokens,
                              double time_to_first_token, double tokens_per_second) {
    ModelTelemetryIdentity identity;
    {
        std::lock_guard<std::mutex> lock(load_mutex_);
        WrappedServer* server = model_name.empty()
            ? get_most_recent_server()
            : find_server_by_model_name(resolve_model_name(model_name));
        identity = get_telemetry_identity(server);
    }
    record_telemetry_for_model(identity, input_tokens, output_tokens,
                               time_to_first_token, tokens_per_second);
}

void Router::update_prompt_tokens(const std::string& model_name, int prompt_tokens) {
    ModelTelemetryIdentity identity;
    {
        std::lock_guard<std::mutex> lock(load_mutex_);
        WrappedServer* server = model_name.empty()
            ? get_most_recent_server()
            : find_server_by_model_name(resolve_model_name(model_name));
        identity = get_telemetry_identity(server);
    }
    record_prompt_tokens_for_model(identity, prompt_tokens);
}

void Router::chat_completion_stream(const std::string& request_body, httplib::DataSink& sink) {
    execute_streaming(request_body, sink, [&](WrappedServer* server) {
        ModelTelemetryIdentity identity = get_telemetry_identity(server);
        server->forward_streaming_request("/v1/chat/completions", request_body, sink, true, 0,
            [this, identity](int input_tokens,
                             int output_tokens,
                             double time_to_first_token,
                             double tokens_per_second) {
                record_telemetry_for_model(identity, input_tokens, output_tokens,
                                           time_to_first_token, tokens_per_second);
                record_prompt_tokens_for_model(identity, input_tokens);
            });
    });
}

void Router::completion_stream(const std::string& request_body, httplib::DataSink& sink) {
    execute_streaming(request_body, sink, [&](WrappedServer* server) {
        ModelTelemetryIdentity identity = get_telemetry_identity(server);
        server->forward_streaming_request("/v1/completions", request_body, sink, true, 0,
            [this, identity](int input_tokens,
                             int output_tokens,
                             double time_to_first_token,
                             double tokens_per_second) {
                record_telemetry_for_model(identity, input_tokens, output_tokens,
                                           time_to_first_token, tokens_per_second);
                record_prompt_tokens_for_model(identity, input_tokens);
            });
    });
}

void Router::responses_stream(const std::string& request_body, httplib::DataSink& sink) {
    execute_streaming(request_body, sink, [&](WrappedServer* server) {
        ModelTelemetryIdentity identity = get_telemetry_identity(server);
        server->forward_streaming_request("/v1/responses", request_body, sink, true, 0,
            [this, identity](int input_tokens,
                             int output_tokens,
                             double time_to_first_token,
                             double tokens_per_second) {
                record_telemetry_for_model(identity, input_tokens, output_tokens,
                                           time_to_first_token, tokens_per_second);
                record_prompt_tokens_for_model(identity, input_tokens);
            });
    });
}

int Router::count_pinned_servers_by_type(ModelType type) const {
    int count = 0;
    for (const auto& server : loaded_servers_) {
        if (server->get_recipe_options().get_recipe() == "cloud") {
            continue;
        }
        if (server->is_backend_alive() && server->get_model_type() == type && server->is_pinned()) {
            count++;
        }
    }
    return count;
}

json Router::get_pinned_model_counts() const {
    std::lock_guard<std::mutex> lock(load_mutex_);
    return {
        {"llm", count_pinned_servers_by_type(ModelType::LLM)},
        {"embedding", count_pinned_servers_by_type(ModelType::EMBEDDING)},
        {"reranking", count_pinned_servers_by_type(ModelType::RERANKING)},
        {"transcription", count_pinned_servers_by_type(ModelType::TRANSCRIPTION)},
        {"image", count_pinned_servers_by_type(ModelType::IMAGE)},
        {"tts", count_pinned_servers_by_type(ModelType::TTS)}
    };
}

void Router::set_model_pinned(const std::string& model_name, bool pinned) {
    std::lock_guard<std::mutex> lock(load_mutex_);
    WrappedServer* server = find_server_by_model_name(model_name);
    if (!server) {
        throw std::runtime_error("Model not loaded: " + model_name);
    }
    server->set_pinned(pinned);
}

} // namespace lemon
