#pragma once

#include <stdexcept>
#include <string>
#include <map>
#include <vector>
#include <mutex>
#include <functional>
#include <nlohmann/json.hpp>
#include "model_types.h"
#include "recipe_options.h"

namespace lemon {

using json = nlohmann::json;

// Thrown by ModelManager::download_model when a pull request names a model
// that (a) is not registered, (b) is not in the filtered-out registry, and
// (c) lacks the `user.` prefix that would make it a new-model registration
// attempt.
//
// CONTRACT: the /pull HTTP handler catches this type and attaches
// {"code": kUnknownModelErrorCode, ...} to the error response. The lemonade
// CLI keys off that code to replace the message with a friendlier one that
// points at `lemonade list` and `lemonade pull CHECKPOINT`. The CLI inlines
// the "unknown_model" literal to avoid pulling this server header into the
// CLI; update cli/lemonade_client.cpp in lockstep if this constant changes.
class UnknownModelError : public std::runtime_error {
public:
    using std::runtime_error::runtime_error;
};
constexpr const char* kUnknownModelErrorCode = "unknown_model";

// Progress information for download operations
struct DownloadProgress {
    std::string file;           // Current file being downloaded
    int file_index = 0;         // Current file index (1-based)
    int total_files = 0;        // Total number of files to download
    size_t bytes_downloaded = 0; // Bytes downloaded for current file
    size_t bytes_total = 0;     // Total bytes for current file
    size_t total_download_size = 0; // Total bytes across ALL files in this download
    size_t bytes_previously_downloaded = 0; // Bytes already on disk (resume offset or skipped file)
    int percent = 0;            // Overall percentage (0-100)
    bool complete = false;      // True when all downloads finished
    std::string error;          // Error message if failed
};

// Callback for download progress updates
// Returns bool: true = continue download, false = cancel download
using DownloadProgressCallback = std::function<bool(const DownloadProgress&)>;

// Image generation defaults for SD models
struct ImageDefaults {
    int steps = 20;
    float cfg_scale = 7.0f;
    int width = 512;
    int height = 512;
    std::string sampling_method;
    float flow_shift = 0.0f;

    bool has_defaults = false;  // True if explicit defaults were provided in JSON
};

struct ModelInfo {
    std::string model_name;
    std::map<std::string, std::string> checkpoints;
    std::map<std::string, std::string> resolved_paths; // Absolute path to model file/directory on disk
    std::string recipe;
    std::vector<std::string> labels;
    std::vector<std::string> composite_models;
    bool suggested = false;
    std::string source;  // "local_upload" for locally uploaded models
    bool downloaded = false;     // Whether model is downloaded and available
    // When true, LlamaCppServer launches llama-server with `-hf <checkpoint>`
    // instead of `-m <gguf> [--mmproj <mmproj>]`. Required for models like
    // Qwen2.5-Omni where llama-server's manual-load path rejects audio content
    // parts — the -hf path drives the dual-clip (vision+audio) context correctly.
    bool hf_load = false;
    double size = 0.0;   // Model size in GB
    int64_t max_context_window = 0;  // Static model-supported text context, when known
    RecipeOptions recipe_options;

    // Multi-model support fields
    ModelType type = ModelType::LLM;      // Model type for LRU cache management
    DeviceType device = DEVICE_NONE;      // Target device(s) for this model

    // Image generation defaults (for sd-cpp models)
    ImageDefaults image_defaults;

    // Utility
    std::string checkpoint(const std::string& type = "main") const { return checkpoints.count(type) ? checkpoints.at(type) : ""; }
    std::string resolved_path(const std::string& type = "main") const { return resolved_paths.count(type) ? resolved_paths.at(type) : ""; }

    std::string mmproj() const { return checkpoint("mmproj"); }
};

class ModelManager {
public:
    explicit ModelManager(const std::string& extra_models_dir = "");

    // Invalidate the models cache (e.g. after backend install/uninstall)
    void invalidate_models_cache();

    // Get all supported models from server_models.json
    std::map<std::string, ModelInfo> get_supported_models();

    // Get downloaded models
    std::map<std::string, ModelInfo> get_downloaded_models();

    // Filter models by available backends
    std::map<std::string, ModelInfo> filter_models_by_backend(
        const std::map<std::string, ModelInfo>& models);

    // Register a user model
    void register_user_model(const std::string& model_name,
                            const json& model_data,
                            const std::string& source = "");

    // Register (if needed) and download a model
    void download_model(const std::string& model_name,
                       const json& model_data,
                       bool do_not_upgrade = false,
                       DownloadProgressCallback progress_callback = nullptr);

    // Download a model
    void download_registered_model(const ModelInfo& info,
                                bool do_not_upgrade = false,
                                DownloadProgressCallback progress_callback = nullptr);

    // Delete a model
    void delete_model(const std::string& model_name);

    // Clean up orphaned files from multi-repo models downloaded in old layout
    nlohmann::json cleanup_orphaned_cache(bool dry_run);

    // Get model info by name
    ModelInfo get_model_info(const std::string& model_name);

    // Resolve a public model reference to its canonical internal name.
    std::string resolve_model_name(const std::string& model_name);

    // Get the public name exposed by Lemonade APIs for a canonical model name.
    std::string get_public_model_name(const std::string& model_name);

    // Check if model exists (in filtered list based on system capabilities)
    bool model_exists(const std::string& model_name);

    // Check if model exists in the raw registry (before filtering)
    // Returns true even for NPU models on systems without NPU
    bool model_exists_unfiltered(const std::string& model_name);

    // Get model info from raw registry (without filtering)
    // Useful for generating helpful error messages about unsupported models
    ModelInfo get_model_info_unfiltered(const std::string& model_name);

    // Get the reason why a model was filtered out (empty string if not filtered)
    // Returns a user-friendly message explaining why the model is not available
    std::string get_model_filter_reason(const std::string& model_name);

    // Check if model is downloaded
    bool is_model_downloaded(const std::string& model_name);

    // Get list of installed FLM models (for caching)
    std::vector<std::string> get_flm_installed_models();

    // Get list of all available FLM models from 'flm list --json'
    std::vector<ModelInfo> get_flm_available_models();

    // Get HuggingFace cache directory (respects HF_HUB_CACHE, HF_HOME, and platform defaults)
    std::string get_hf_cache_dir() const;

    // Set extra models directory for GGUF discovery
    void set_extra_models_dir(const std::string& dir);

    void save_model_options(const ModelInfo& info);

private:
    json load_server_models();
    json load_optional_json(const std::string& path);
    void save_user_models(const json& user_models);

    std::string get_user_models_file();
    std::string get_recipe_options_file();

    // Cache management
    void build_cache();
    void add_model_to_cache(const std::string& model_name);
    void update_model_options_in_cache(const ModelInfo& info);
    void update_model_in_cache(const std::string& model_name, bool downloaded);
    void remove_model_from_cache(const std::string& model_name);

    // Resolve model checkpoint to absolute path on disk
    std::string resolve_model_path(const ModelInfo& info, const std::string& type, const std::string& checkpoint) const;
    void resolve_all_model_paths(ModelInfo& info);

    // Download from a JSON manifest
    void download_from_manifest(const json& manifest, std::map<std::string, std::string>& headers, DownloadProgressCallback progress_callback);

    // Download from Hugging Face
    void download_from_huggingface(const ModelInfo& info,
                                   DownloadProgressCallback progress_callback = nullptr);

    // Download from FLM
    void download_from_flm(const std::string& checkpoint,
                          bool do_not_upgrade = true,
                          DownloadProgressCallback progress_callback = nullptr);

    // Discover GGUF models from extra_models_dir
    std::map<std::string, ModelInfo> discover_extra_models() const;

    json server_models_;
    json user_models_;
    json recipe_options_;
    std::string extra_models_dir_;  // Secondary directory for GGUF model discovery

    // Cache of all models with their download status
    mutable std::mutex models_cache_mutex_;
    mutable std::map<std::string, ModelInfo> models_cache_;
    mutable std::map<std::string, std::string> public_model_aliases_;  // public name -> canonical name
    mutable std::map<std::string, std::string> canonical_public_names_;  // canonical name -> public name
    mutable std::map<std::string, std::string> filtered_out_models_;  // model_name -> filter reason
    mutable bool cache_valid_ = false;

    void rebuild_public_model_aliases_locked();
};

} // namespace lemon
