#include "cli-context.h"
#include "cli-ui.h"

#include "arg.h"
#include "base64.hpp"
#include "log.h"
#include "console.h"

#define JSON_ASSERT GGML_ASSERT
#include <nlohmann/json.hpp>

#include <algorithm>
#include <cctype>
#include <filesystem>
#include <fstream>
#include <map>
#include <set>

using json = nlohmann::ordered_json;

struct cli_context_impl {
    json messages      = json::array();
    json pending_media = json::array(); // staged multimodal content parts
};

cli_context::cli_context(const common_params & params) : params(params), impl(new cli_context_impl()) {}

cli_context::~cli_context() {
    shutdown();
}

std::atomic<bool> & cli_context::interrupted() {
    static std::atomic<bool> flag = false;
    return flag;
}

static bool should_stop() {
    return cli_context::interrupted().load();
}

static constexpr size_t FILE_GLOB_MAX_RESULTS = 100;

const char * LLAMA_ASCII_LOGO = R"(
▄▄ ▄▄
██ ██
██ ██  ▀▀█▄ ███▄███▄  ▀▀█▄    ▄████ ████▄ ████▄
██ ██ ▄█▀██ ██ ██ ██ ▄█▀██    ██    ██ ██ ██ ██
██ ██ ▀█▄██ ██ ██ ██ ▀█▄██ ██ ▀████ ████▀ ████▀
                                    ██    ██
                                    ▀▀    ▀▀
)";

// number of values an arg consumes on the command line
static int arg_num_values(const common_arg & opt) {
    if (opt.value_hint_2 != nullptr) {
        return 2;
    }
    if (opt.value_hint != nullptr) {
        return 1;
    }
    return 0;
}

static std::string format_error_message(const json & err) {
    if (err.contains("error") && err.at("error").is_object()) {
        const auto & e = err.at("error");
        if (e.contains("message") && e.at("message").is_string()) {
            return e.at("message").get<std::string>();
        }
    }
    return err.dump();
}

// err is the raw response body of a failed request; it may or may not be JSON
static std::string format_error_message(const std::string & err) {
    json parsed = json::parse(err, nullptr, false);
    if (!parsed.is_discarded()) {
        return format_error_message(parsed);
    }
    return err;
}

static std::string media_type_from_ext(const std::string & fname) {
    std::string ext = std::filesystem::path(fname).extension().string();
    std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) { return std::tolower(c); });
    if (ext == ".wav" || ext == ".mp3") {
        return "audio";
    }
    if (ext == ".mp4" || ext == ".avi" || ext == ".mkv" || ext == ".mov" || ext == ".webm") {
        return "video";
    }
    return "image";
}

bool cli_context::init() {
    ui::init(params);

    std::optional<ui::spinner> spinner;

    bool use_external_server = !params.server_base.empty();
    if (use_external_server) {
        std::string base = params.server_base;
        while (!base.empty() && base.back() == '/') {
            base.pop_back();
        }
        client.server_base = base;

        spinner.emplace("Connecting to server at " + base);
    } else {
        if (params.model.path.empty() && params.model.url.empty() &&
                params.model.hf_repo.empty() && params.model.docker_repo.empty()) {
            ui::show_error(
                "no model specified",
                "use -m <file.gguf> or -hf <user/repo> to run a local model,\n"
                "or --server-base <url> to connect to a running llama-server"
            );
            return false;
        }

        spinner.emplace("\n\nLoading model...");

        server.emplace();
        if (!server->start(params)) {
            ui::show_error("server start failed");
            return false;
        }
        if (!server->wait_ready(should_stop)) {
            if (!should_stop()) {
                ui::show_error("the server exited before becoming ready");
            }
            return false;
        }
        client.server_base = server->address();
    }

    // for --server-base this is the main availability check; for a spawned
    // server it is a cheap sanity check on top of the ready signal
    auto is_aborted = [this]() {
        return should_stop() || (server && !server->alive());
    };
    bool healthy = false;
    try {
        healthy = client.wait_health(is_aborted);
    } catch (const std::exception & e) {
        client.last_error = e.what();
    }
    if (!healthy) {
        if (!should_stop()) {
            ui::show_error(client.last_error);
        }
        return false;
    }

    if (use_external_server) {
        spinner.reset();
        try {
            if (!list_and_ask_models()) {
                return false;
            }
        } catch (const json::parse_error & e) {
            ui::show_error(e.what());
            ui::show_message("This might be caused by an incorrect server-base endpoint URL");
            return false;
        } catch (const std::exception & e) {
            ui::show_error(e.what());
            return false;
        }

        // restore the spinner for the next step
        spinner.emplace("Waiting for server...");
    }

    fetch_server_props();

    if (!params.out_file.empty()) {
        output_file.emplace(params.out_file);
        if (!output_file->is_open()) {
            ui::show_error(string_format("failed to open output file '%s'", params.out_file.c_str()));
            return false;
        }
    }

    return true;
}

void cli_context::fetch_server_props() {
    try {
        json props = json::parse(client.get("/props"));
        model_name = props.value("model_alias", "");
        if (model_name.empty()) {
            const std::string path = props.value("model_path", "");
            if (!path.empty()) {
                model_name = std::filesystem::path(path).filename().string();
            }
        }
        model_ftype = props.value("model_ftype", "");
        build_info = props.value("build_info", "");
        if (props.contains("modalities") && props.at("modalities").is_object()) {
            const auto & modalities = props.at("modalities");
            has_vision = modalities.value("vision", false);
            has_audio  = modalities.value("audio", false);
            has_video  = modalities.value("video", false);
        }
    } catch (const std::exception & e) {
        // /props can be disabled on remote servers; not fatal
        LOG_DBG("failed to fetch /props: %s\n", e.what());
    }
}

bool cli_context::list_and_ask_models() {
    json resp = json::parse(client.get("/v1/models"));
    if (!resp.contains("data") || !resp.at("data").is_array()) {
        throw std::runtime_error("invalid response from /v1/models");
    }
    std::vector<std::string> models;
    std::vector<std::string> models_display;
    for (const auto & m : resp.at("data")) {
        if (!m.contains("id") || !m.at("id").is_string()) {
            continue;
        }
        std::string name = m.at("id").get<std::string>();
        std::string display = name;
        if (m.contains("aliases") && m.at("aliases").is_array()) {
            std::vector<std::string> aliases;
            for (const auto & a : m.at("aliases")) {
                if (a.is_string()) {
                    aliases.push_back(a.get<std::string>());
                }
            }
            if (!aliases.empty()) {
                display += " (" + string_join(aliases, ", ") + ")";
            }
        }
        models.push_back(name);
        models_display.push_back(display);
    }

    // only one model: use it without asking
    if (models.size() == 1) {
        model_name = models[0];
        client.model = model_name;
        return true;
    }

    std::string message = "\nAvailable models:";
    for (size_t i = 0; i < models_display.size(); ++i) {
        message += "\n  " + std::to_string(i + 1) + ". " + models_display[i];
    }
    message += "\n";
    ui::show_message(message);
    std::string selection;
    while (selection.empty()) {
        if (should_stop()) {
            return false;
        }
        ui::user_turn user_turn;
        selection = user_turn.read_input(false, "Select model by number: ");
        if (selection.empty()) {
            continue;
        }
        try {
            size_t idx = std::stoul(selection);
            if (idx > 0 && idx <= models.size()) {
                model_name = models[idx - 1];
                client.model = model_name;
                ui::show_message("Selected model: " + model_name);
                break;
            }
        } catch (...) {
            // ignore
        }
        ui::show_error("Invalid selection. Please enter a valid number.");
        selection.clear();
        continue;
    }
    return true;
}

void cli_context::add_system_prompt() {
    if (!params.system_prompt.empty()) {
        impl->messages.push_back({
            {"role",    "system"},
            {"content", params.system_prompt}
        });
    }
}

void cli_context::push_user_message(const std::string & text) {
    json content;
    if (impl->pending_media.empty()) {
        content = text;
    } else {
        // multimodal message: media parts first, then the text
        content = impl->pending_media;
        content.push_back({
            {"type", "text"},
            {"text", text}
        });
        impl->pending_media = json::array();
    }
    impl->messages.push_back({
        {"role",    "user"},
        {"content", content}
    });
}

bool cli_context::stage_media_file(const std::string & fname, const std::string & type) {
    std::ifstream file(fname, std::ios::binary);
    if (!file) {
        return false;
    }
    std::string data((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
    std::string encoded = base64::encode(data);

    if (type == "audio") {
        std::string ext = std::filesystem::path(fname).extension().string();
        std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) { return std::tolower(c); });
        impl->pending_media.push_back({
            {"type", "input_audio"},
            {"input_audio", {
                {"data",   encoded},
                {"format", ext == ".mp3" ? "mp3" : "wav"}
            }}
        });
    } else if (type == "video") {
        impl->pending_media.push_back({
            {"type", "input_video"},
            {"input_video", {
                {"data", encoded}
            }}
        });
    } else {
        // the server detects the actual image type from the data
        impl->pending_media.push_back({
            {"type", "image_url"},
            {"image_url", {
                {"url", "data:image/unknown;base64," + encoded}
            }}
        });
    }
    return true;
}

void cli_context::write_output_file(const std::string & content) {
    if (output_file) {
        (*output_file) << content;
        output_file->flush();
    }
}

bool cli_context::generate_completion(generated_content & content_out, cli_timings & timings) {
    json body = {
        {"messages",          impl->messages},
        {"stream",            true},
        // in order to get timings even when we cancel mid-way
        {"timings_per_token", true},
    };
    if (!client.model.empty()) {
        body["model"] = client.model;
    }

    bool stream_error = false;

    ui::assistant_turn a;

    std::string err = client.post_sse("/v1/chat/completions", body.dump(), should_stop, [&](const std::string & payload) {
        json chunk = json::parse(payload, nullptr, false);
        if (chunk.is_discarded()) {
            return;
        }
        if (chunk.contains("error")) {
            stream_error = true;
            ui::show_error(format_error_message(chunk));
            return;
        }
        if (chunk.contains("timings")) {
            const auto & t = chunk.at("timings");
            timings.prompt_per_second    = t.value("prompt_per_second",    0.0);
            timings.predicted_per_second = t.value("predicted_per_second", 0.0);
        }
        if (!chunk.contains("choices") || !chunk.at("choices").is_array() || chunk.at("choices").empty()) {
            return;
        }
        const auto & choice = chunk.at("choices").at(0);
        if (!choice.contains("delta")) {
            return;
        }
        const auto & delta = choice.at("delta");
        if (delta.contains("reasoning_content") && delta.at("reasoning_content").is_string()) {
            const std::string text = delta.at("reasoning_content").get<std::string>();
            if (!text.empty()) {
                content_out.reasoning += text;
                a.push(ui::ASSISTANT_DISPLAY_MODE_REASONING, text);
            }
        }
        if (delta.contains("content") && delta.at("content").is_string()) {
            const std::string text = delta.at("content").get<std::string>();
            if (!text.empty()) {
                content_out.content += text;
                a.push(ui::ASSISTANT_DISPLAY_MODE_CONTENT, text);
            }
        }
    });

    cli_context::interrupted().store(false);

    if (!err.empty()) {
        ui::show_error(format_error_message(err));
        return false;
    }
    return !stream_error;
}

int cli_context::run() {
    add_system_prompt();

    std::string modalities = "text";
    if (has_vision) {
        modalities += ", vision";
    }
    if (has_audio) {
        modalities += ", audio";
    }
    if (has_video) {
        modalities += ", video";
    }

    std::string banner;
    banner += "\n";
    banner += LLAMA_ASCII_LOGO;
    banner += "\n";
    banner += "build      : " + build_info + "\n";
    banner += "model      : " + model_name + "\n";
    if (!model_ftype.empty()) {
        banner += "ftype      : " + model_ftype + "\n";
    }
    banner += "modalities : " + modalities + "\n";
    if (!params.system_prompt.empty()) {
        banner += "using custom system prompt\n";
    }
    banner += "\n";
    banner += "available commands:\n";
    banner += "  /exit or Ctrl+C     stop or exit\n";
    banner += "  /regen              regenerate the last response\n";
    banner += "  /clear              clear the chat history\n";
    banner += "  /read <file>        add a text file\n";
    banner += "  /glob <pattern>     add text files using globbing pattern\n";
    if (has_vision) {
        banner += "  /image <file>       add an image file\n";
    }
    if (has_audio) {
        banner += "  /audio <file>       add an audio file\n";
    }
    if (has_video) {
        banner += "  /video <file>       add a video file\n";
    }
    banner += "\n";

    ui::show_message(banner);

    // interactive loop
    std::string cur_msg;

    auto add_text_file = [&](const std::string & fname) -> bool {
        std::ifstream file(fname, std::ios::binary);
        if (!file) {
            ui::show_error(string_format("file does not exist or cannot be opened: '%s'", fname.c_str()));
            return false;
        }
        std::string content((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
        cur_msg += "--- File: ";
        cur_msg += fname;
        cur_msg += " ---\n";
        cur_msg += content;
        ui::show_message(string_format("Loaded text from '%s'", fname.c_str()));
        return true;
    };

    while (true) {
        std::string buffer;
        {
            ui::user_turn user_turn;

            if (params.prompt.empty()) {
                buffer = user_turn.read_input(params.multiline_input);
            } else {
                // process input prompt from args
                for (auto & fname : params.image) {
                    if (!stage_media_file(fname, media_type_from_ext(fname))) {
                        ui::show_error(string_format("file does not exist or cannot be opened: '%s'", fname.c_str()));
                        break;
                    }
                    ui::show_message(string_format("Loaded media from '%s'", fname.c_str()));
                }
                buffer = params.prompt;
                user_turn.echo(buffer);
                params.prompt.clear(); // only use it once
            }
        }

        if (should_stop()) {
            cli_context::interrupted().store(false);
            break;
        }

        // remove trailing newline
        if (!buffer.empty() && buffer.back() == '\n') {
            buffer.pop_back();
        }

        // skip empty messages
        if (buffer.empty()) {
            continue;
        }

        bool add_user_msg = true;

        // process commands
        if (string_starts_with(buffer, "/exit")) {
            break;
        } else if (string_starts_with(buffer, "/regen")) {
            if (impl->messages.size() >= 2) {
                size_t last_idx = impl->messages.size() - 1;
                impl->messages.erase(last_idx);
                add_user_msg = false;
            } else {
                ui::show_error("No message to regenerate.");
                continue;
            }
        } else if (string_starts_with(buffer, "/clear")) {
            impl->messages.clear();
            add_system_prompt();

            impl->pending_media = json::array();
            ui::show_message("Chat history cleared.");
            continue;
        } else if (
                (string_starts_with(buffer, "/image ") && has_vision) ||
                (string_starts_with(buffer, "/audio ") && has_audio) ||
                (string_starts_with(buffer, "/video ") && has_video)) {
            std::string type = buffer.substr(1, 5);
            // just in case (bad copy-paste for example), we strip all trailing/leading spaces
            std::string fname = string_strip(buffer.substr(7));
            if (!stage_media_file(fname, type)) {
                ui::show_error(string_format("file does not exist or cannot be opened: '%s'", fname.c_str()));
                continue;
            }
            ui::show_message(string_format("Loaded media from '%s'", fname.c_str()));
            write_output_file(string_format("User: Added media: %s\n", fname.c_str()));
            continue;
        } else if (string_starts_with(buffer, "/read ")) {
            std::string fname = string_strip(buffer.substr(6));
            add_text_file(fname);
            write_output_file(string_format("User: Added text file: %s\n", fname.c_str()));
            continue;
        } else if (string_starts_with(buffer, "/glob ")) {
            std::error_code ec;
            size_t count = 0;
            auto curdir = std::filesystem::current_path();
            std::string pattern = string_strip(buffer.substr(6));
            std::filesystem::path rel_path;

            auto startglob = pattern.find_first_of("![*?");
            if (startglob != std::string::npos && startglob != 0) {
                auto endpath = pattern.substr(0, startglob).find_last_of('/');
                if (endpath != std::string::npos) {
                    std::string rel_pattern = pattern.substr(0, endpath);
#if !defined(_WIN32)
                    if (string_starts_with(rel_pattern, '~')) {
                        const char * home = std::getenv("HOME");
                        if (home && home[0]) {
                            rel_pattern = home + rel_pattern.substr(1);
                        }
                    }
#endif
                    rel_path = rel_pattern;
                    pattern.erase(0, endpath + 1);
                    curdir /= rel_path;
                }
            }

            for (const auto & entry : std::filesystem::recursive_directory_iterator(curdir,
                    std::filesystem::directory_options::skip_permission_denied, ec)) {
                if (!entry.is_regular_file()) {
                    continue;
                }

                std::string rel = std::filesystem::relative(entry.path(), curdir, ec).string();
                if (ec) {
                    ec.clear();
                    continue;
                }
                std::replace(rel.begin(), rel.end(), '\\', '/');

                if (!glob_match(pattern, rel)) {
                    continue;
                }

                const std::string full_path = (curdir / rel).string();
                if (!add_text_file(full_path)) {
                    continue;
                }
                write_output_file(string_format("User: Added text file: %s\n", full_path.c_str()));

                if (++count >= FILE_GLOB_MAX_RESULTS) {
                    ui::show_error(string_format("Maximum number of globbed files allowed (%zu) reached.", FILE_GLOB_MAX_RESULTS));
                    break;
                }
            }
            continue;
        } else {
            // not a command
            cur_msg += buffer;
        }

        // generate response
        if (add_user_msg) {
            push_user_message(cur_msg);
            write_output_file(string_format("User:\n%s\n\n", cur_msg.c_str()));
            cur_msg.clear();
        }

        cli_timings timings;
        generated_content content;
        generate_completion(content, timings);

        json assistant_msg = {
            {"role",    "assistant"},
            {"content", content.content}
        };
        if (!content.reasoning.empty()) {
            assistant_msg["reasoning_content"] = content.reasoning;
        }
        impl->messages.push_back(std::move(assistant_msg));

        if (output_file) {
            std::string out_content = "Assistant:\n";
            if (!content.reasoning.empty()) {
                out_content += "[Start thinking]\n\n";
                out_content += content.reasoning;
                out_content += "[End thinking]\n\n";
            }
            out_content += content.content;
            if (!out_content.empty() && out_content.back() != '\n') {
                out_content += "\n";
            }
            out_content += "\n";
            write_output_file(out_content);
        }

        if (params.show_timings) {
            ui::show_info(string_format(
                "\n[ Prompt: %.1f t/s | Generation: %.1f t/s ]",
                timings.prompt_per_second,
                timings.predicted_per_second
            ));
        }

        if (params.single_turn) {
            break;
        }
    }

    ui::show_message("\n\nExiting...");

    return 0;
}

void cli_context::shutdown() {
    if (server) {
        server->stop();
        server.reset();
    }
    if (output_file) {
        output_file->close();
        output_file.reset();
    }
}