#include <lemon/recipe_options.h>
#include <lemon/utils/custom_args.h>
#include <nlohmann/json.hpp>
#include <map>
#ifdef LEMONADE_CLI
#include <CLI/CLI.hpp>
#else
#include <lemon/system_info.h>
#endif

namespace lemon {

using json = nlohmann::json;

static const json DEFAULTS = {
    {"ctx_size", -1},  // -1 triggers auto-resolution (memory + arch metadata)
    {"merge_args", true},
    {"llamacpp_device", ""},
    {"llamacpp_backend", ""},  // Will be overridden dynamically
    {"llamacpp_args", ""},
    {"sd-cpp_backend", ""},   // "" means auto-detect (mapped from "auto" in config.json)
    {"sdcpp_args", ""},
    {"whispercpp_backend", ""},  // "" means auto-detect (mapped from "auto" in config.json)
    {"whispercpp_args", ""},
    {"moonshine_args", ""},      // Custom arguments to pass to moonshine-server
    // Image generation defaults (for sd-cpp recipe)
    // These are recipe-level defaults only, not CLI arguments — per reviewer guidance,
    // there are too many image gen params for CLI flags, and no universal defaults.
    {"steps", 20},
    {"cfg_scale", 7.0},
    {"width", 512},
    {"height", 512},
    {"sampling_method", ""},
    {"flow_shift", 0.0},
    // vLLM-specific options
    {"vllm_backend", ""},  // "" means auto-detect
    {"vllm_args", ""},     // Custom arguments to pass to vllm-server
    // Cloud recipe has no backend variants (provider selection lives on the
    // per-model cloud_provider field). The empty string satisfies Router's
    // per-backend-args lookup; cloud reads no backend-specific config.
    {"cloud_backend", ""},

    // Auto-eviction options
    {"auto_evict", nullptr},          // nullptr means fallback to global config
    {"evict_idle_timeout", 300},      // Default hard idle timeout (5 mins)
    {"downsize_idle_timeout", 60},    // Default soft idle timeout (1 min)
    {"evict_weight_factor", 1.0},     // Eviction-protection weight (higher = more protected)
    {"pinned", false}
};


// Mapping from flat option names to CLI flags (used by to_cli_options)
// Note: Image generation params (steps, cfg_scale, width, height, sampling_method,
// flow_shift) are recipe-level defaults only — not exposed as CLI arguments.
// Runtime options (diffusion_fa, offload_to_cpu) go through --sdcpp-args.
static const std::map<std::string, std::string> OPTION_TO_CLI_FLAG = {
    {"ctx_size", "--ctx-size"},
    {"merge_args", "--merge-args"},
    {"llamacpp_backend", "--llamacpp"},
    {"llamacpp_device", "--llamacpp-device"},
    {"llamacpp_args", "--llamacpp-args"},
    {"sd-cpp_backend", "--sdcpp"},
    {"sdcpp_args", "--sdcpp-args"},
    {"whispercpp_backend", "--whispercpp"},
    {"whispercpp_args", "--whispercpp-args"},
    {"moonshine_args", "--moonshine-args"},
    {"vllm_backend", "--vllm"},
    {"vllm_args", "--vllm-args"}
};

static std::vector<std::string> get_keys_for_recipe(const std::string& recipe) {
    std::vector<std::string> keys;
    if (recipe == "llamacpp") {
        keys = {"ctx_size", "llamacpp_device", "llamacpp_backend", "llamacpp_args", "merge_args"};
    } else if (recipe == "whispercpp") {
        keys = {"whispercpp_backend", "whispercpp_args", "merge_args"};
    } else if (recipe == "moonshine") {
        keys = {"moonshine_args", "merge_args"};
    } else if (recipe == "flm") {
        return {"ctx_size", "merge_args"};
    } else if (recipe == "ryzenai-llm") {
        keys = {"ctx_size"};
    } else if (recipe == "sd-cpp") {
        keys = {"sd-cpp_backend", "sdcpp_args", "steps", "cfg_scale", "width", "height", "sampling_method", "flow_shift", "merge_args"};
    } else if (recipe == "vllm") {
        keys = {"ctx_size", "vllm_backend", "vllm_args", "merge_args"};
    }

    // Add auto-eviction options for all recipes
    keys.push_back("auto_evict");
    keys.push_back("evict_idle_timeout");
    keys.push_back("downsize_idle_timeout");
    keys.push_back("evict_weight_factor");
    keys.push_back("pinned");

    return keys;
}

static bool is_empty_option(json option) {
    return option.is_null() ||
           (option.is_number() && (option == -1)) ||
           (option.is_string() && (option == "" || option == "auto"));
}


#ifndef LEMONADE_CLI
static bool try_get_backend_options(const std::string& opt_name, SystemInfo::SupportedBackendsResult& result) {
    // Generic handling for any *_backend option
    // Pattern: {recipe}_backend -> get supported backends for {recipe}
    const std::string backend_suffix = "_backend";
    bool is_backend_option = opt_name.size() > backend_suffix.size() &&
        opt_name.compare(opt_name.size() - backend_suffix.size(), backend_suffix.size(), backend_suffix) == 0;

    if (is_backend_option) {
        // Extract recipe name (everything before "_backend")
        std::string recipe = opt_name.substr(0, opt_name.size() - backend_suffix.size());
        auto tmp = SystemInfo::get_supported_backends(recipe);
        result.backends = tmp.backends;
    }

    return is_backend_option;
}
#endif

std::vector<std::string> RecipeOptions::to_cli_options(const json& raw_options) {
    std::vector<std::string> cli;

    for (auto& [opt_name, cli_flag] : OPTION_TO_CLI_FLAG) {
        if (raw_options.contains(opt_name)) {
            auto val = raw_options[opt_name];
            if (!val.is_null() && val != "") {
                cli.push_back(cli_flag);
                if (val.is_number_float()) {
                    cli.push_back(std::to_string((double) val));
                } else if (val.is_number_integer()) {
                    cli.push_back(std::to_string((int) val));
                } else {
                    cli.push_back(val);
                }
            }
        }
    }

    return cli;
}

std::vector<std::string> RecipeOptions::known_keys() {
    std::vector<std::string> keys;
    for (auto& [key, value] : DEFAULTS.items()) {
        keys.push_back(key);
    }
    return keys;
}

RecipeOptions::RecipeOptions(const std::string& recipe, const json& options) {
    recipe_ = recipe;
    std::vector<std::string> to_copy = get_keys_for_recipe(recipe_);

    for (auto key : to_copy) {
        if (options.contains(key) && !is_empty_option(options[key])) {
            options_[key] = options[key];
        }
    }
}

static std::string format_option_for_logging(const json& opt) {
    if (opt.is_null() || opt == "") return "(none)";
    if (opt.is_boolean()) return opt.get<bool>() ? "true" : "false";
    if (opt.is_number_float()) return std::to_string((double) opt);
    if (opt.is_number_integer()) return std::to_string((int) opt);
    return opt;
}

json RecipeOptions::to_json() const {
    return options_;
}

std::string RecipeOptions::to_log_string(bool resolve_defaults) const {
    std::vector<std::string> to_log = get_keys_for_recipe(recipe_);
    std::string log_string = "";
    bool first = true;
    for (auto key : to_log) {
        if (resolve_defaults || options_.contains(key)) {
            if (!first) log_string += ", ";
            first = false;
            log_string += key + "=" + format_option_for_logging(get_option(key));
        }
    }

    return log_string;
}

RecipeOptions RecipeOptions::inherit(const RecipeOptions& options) const {
    json merged = options_;
    bool merge_args = options_.contains("merge_args") ? options_["merge_args"].get<bool>() : options.get_option("merge_args").get<bool>();

    for (auto it = options.options_.begin(); it != options.options_.end(); ++it) {
        if (merge_args && it.key().size() >= 5 && it.key().substr(it.key().size() - 5) == "_args") {
            // Special handling for _args options: parse, merge maps, re-stringify
            std::string target_str = "";
            if (merged.contains(it.key()) && merged[it.key()].is_string()) {
                target_str = merged[it.key()];
            }

            std::string incoming_str = it.value().is_string() ? it.value() : "";

            if (target_str.empty()) {
                merged[it.key()] = incoming_str;
            } else if (incoming_str.empty()) {
                merged[it.key()] = target_str;
            } else {
                auto target_tokens = lemon::utils::parse_custom_args(target_str, true);
                auto incoming_tokens = lemon::utils::parse_custom_args(incoming_str, true);

                auto target_map = lemon::utils::build_custom_args_map(target_tokens);
                auto incoming_map = lemon::utils::build_custom_args_map(incoming_tokens);

                auto merged_map = lemon::utils::merge_args_maps(target_map, incoming_map);
                merged[it.key()] = lemon::utils::map_to_args_string(merged_map);
            }
        } else if (!merged.contains(it.key()) && !is_empty_option(it.value())) {
            merged[it.key()] = it.value();
        }
    }

    return RecipeOptions(recipe_, merged);
}

json RecipeOptions::get_option(const std::string& opt) const {
    if (options_.contains(opt)) {
        return options_[opt];
    }
#ifndef LEMONADE_CLI
    // Dynamic defaults for backends if not explicitly set
    SystemInfo::SupportedBackendsResult backend_result;
    if (try_get_backend_options(opt, backend_result)) {
        if (!backend_result.backends.empty()) {
            return backend_result.backends[0];
        }
    }
#endif
    return DEFAULTS.contains(opt) ? DEFAULTS[opt] : json();
}

void RecipeOptions::set_option(const std::string& opt, const json& value) {
    options_[opt] = value;
}

#ifdef LEMONADE_CLI
// CLI_OPTIONS used only by the lemonade CLI client for add_cli_options
static const json CLI_OPTIONS = {
    {"--ctx-size", {{"option_name", "ctx_size"}, {"type_name", "SIZE"}, {"help", "Context size for the model"}, {"group", "General Options"}}},
    {"--merge-args", {{"option_name", "merge_args"}, {"type_name", "BOOL"}, {"help", "Merge global and model arguments when loading the model"}, {"group", "General Options"}}},
    {"--llamacpp", {{"option_name", "llamacpp_backend"}, {"type_name", "BACKEND"}, {"help", "LlamaCpp backend to use"}, {"group", "Llama.cpp Backend Options"}}},
    {"--llamacpp-device", {{"option_name", "llamacpp_device"}, {"type_name", "DEVICES"}, {"help", "Comma-separated list of accelerator devices to use (e.g. Vulkan0)"}, {"group", "Llama.cpp Backend Options"}}},
    {"--llamacpp-args", {{"option_name", "llamacpp_args"}, {"type_name", "ARGS"}, {"help", "Custom arguments to pass to llama-server"}, {"group", "Llama.cpp Backend Options"}}},
    {"--sdcpp", {{"option_name", "sd-cpp_backend"}, {"type_name", "BACKEND"}, {"help", "SD.cpp backend to use"}, {"group", "Stable Diffusion Options"}}},
    {"--sdcpp-args", {{"option_name", "sdcpp_args"}, {"type_name", "ARGS"}, {"help", "Custom arguments to pass to sd-server (must not conflict with managed args)"}, {"group", "Stable Diffusion Options"}}},
    {"--whispercpp", {{"option_name", "whispercpp_backend"}, {"type_name", "BACKEND"}, {"help", "WhisperCpp backend to use"}, {"group", "Whisper.cpp Options"}}},
    {"--whispercpp-args", {{"option_name", "whispercpp_args"}, {"type_name", "ARGS"}, {"help", "Custom arguments to pass to whisper-server"}, {"group", "Whisper.cpp Options"}}},
    {"--moonshine-args", {{"option_name", "moonshine_args"}, {"type_name", "ARGS"}, {"help", "Custom arguments to pass to moonshine-server"}}},
    {"--vllm", {{"option_name", "vllm_backend"}, {"type_name", "BACKEND"}, {"help", "vLLM backend to use"}, {"group", "vLLM Options"}}},
    {"--vllm-args", {{"option_name", "vllm_args"}, {"type_name", "ARGS"}, {"help", "Custom arguments to pass to vllm-server"}, {"group", "vLLM Options"}}},
    // Note: Image gen params (--steps, --cfg-scale, --width, --height) removed — recipe-level only.
    // Runtime options (--diffusion-fa, --offload-to-cpu) go through --sdcpp-args.
};

void RecipeOptions::add_cli_options(CLI::App& app, json& storage) {
    for (auto& [key, opt] : CLI_OPTIONS.items()) {
        const std::string opt_name = opt["option_name"];
        CLI::Option* o;
        json defval = DEFAULTS[opt_name];

        if (defval.is_number_float()) {
            o = app.add_option_function<double>(key, [opt_name, &storage = storage](double val) { storage[opt_name] = val; }, opt["help"]);
            o->default_val((double) defval);
        } else if (defval.is_number_integer()) {
            o = app.add_option_function<int>(key, [opt_name, &storage = storage](int val) { storage[opt_name] = val; }, opt["help"]);
            o->default_val((int) defval);
        } else if (defval.is_boolean()) {
            o = app.add_flag_function(key + ",!" + lemon::utils::negate_flag(key), [opt_name, defval, &storage = storage](std::int64_t val) { storage[opt_name] = val == 0 ? defval.get<bool>() : val > 0; }, opt["help"]);
            o->default_val((bool) defval);
        } else {
            o = app.add_option_function<std::string>(key, [opt_name, &storage = storage](const std::string& val) { storage[opt_name] = val; }, opt["help"]);
            o->default_val(defval);
        }

        o->type_name(opt["type_name"]);
        o->group(opt.value("group", "Options"));
    }
}
#endif

}
