dstalk/dstalk-cli/src/main.cpp

// ============================================================================
// dstalk-cli — 命令行前端 (使用插件化架构)
// ============================================================================
// 通过 dstalk_host.h API 初始化核心，然后查询插件服务 vtable 调用功能。
// ============================================================================

#include <algorithm>
#include <atomic>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <string>
#include <system_error>
#include <vector>

#ifdef _WIN32
#include <windows.h>
#include <io.h>
#else
#include <signal.h>
#include <termios.h>
#include <unistd.h>
#endif

#include "dstalk/dstalk_host.h"

// ---- ANSI 简写 ----
#define CLR_RESET   "\033[0m"
#define CLR_CYAN    "\033[36m"
#define CLR_YELLOW  "\033[33m"
#define CLR_GREEN   "\033[32m"
#define CLR_RED     "\033[31m"
#define CLR_DIM     "\033[2m"
#define CLR_BOLD    "\033[1m"

// ---- 退出码 ----
#define EXIT_OK            0
#define EXIT_INIT_FAIL     1
#define EXIT_AI_ERROR      2
#define EXIT_SVC_UNAVAIL   3

// ---- 服务 vtable 指针 ----
static const dstalk_ai_service_t*      g_ai      = nullptr;
static const dstalk_session_service_t* g_session = nullptr;
static const dstalk_file_io_service_t* g_file_io = nullptr;

// ---- 运行时状态 ----
static std::string       g_current_model;
static std::atomic<bool> g_quit_requested{false};

// ---- Ctrl+C 信号处理 ----
#ifdef _WIN32
static BOOL WINAPI on_console_event(DWORD event)
{
    if (event == CTRL_C_EVENT || event == CTRL_BREAK_EVENT) {
        g_quit_requested = true;
        return TRUE;
    }
    return FALSE;
}
#else
static void on_signal(int /*sig*/)
{
    g_quit_requested = true;
}
#endif

// ---- 工具函数 ----
static void print_banner()
{
    std::printf("%sdstalk v0.1.0%s  |  %sdstalk AI%s  |  "
                "%s/help%s 查看帮助  |  %s/quit%s 退出\n",
                CLR_CYAN CLR_BOLD, CLR_RESET,
                CLR_GREEN, CLR_RESET,
                CLR_DIM, CLR_RESET,
                CLR_DIM, CLR_RESET);
}

static void print_help()
{
    std::printf("\n%s命令列表:%s\n", CLR_BOLD, CLR_RESET);
    std::printf("  %s/help%s                  显示此帮助\n", CLR_YELLOW, CLR_RESET);
    std::printf("  %s/quit%s                  退出程序\n", CLR_YELLOW, CLR_RESET);
    std::printf("  %s/clear%s                 清空当前会话上下文\n", CLR_YELLOW, CLR_RESET);
    std::printf("  %s/context%s               显示当前 Token 数和消息条数\n", CLR_YELLOW, CLR_RESET);
    std::printf("  %s/status%s                显示当前运行状态（脱敏）\n", CLR_YELLOW, CLR_RESET);
    std::printf("  %s/model <name>%s          切换模型\n", CLR_YELLOW, CLR_RESET);
    std::printf("  %s/file list [path]%s      列出目录内容\n", CLR_YELLOW, CLR_RESET);
    std::printf("  %s/file show <path>%s      查看文件内容\n", CLR_YELLOW, CLR_RESET);
    std::printf("  %s/file read <path>%s      读取文件内容\n", CLR_YELLOW, CLR_RESET);
    std::printf("  %s/file write <p> <c>%s    写入文件\n", CLR_YELLOW, CLR_RESET);
    std::printf("  %s/save <path>%s           保存会话\n", CLR_YELLOW, CLR_RESET);
    std::printf("  %s/load <path>%s           恢复会话\n", CLR_YELLOW, CLR_RESET);
    std::printf("\n直接输入问题即可与 AI 对话。\n\n");
}

static void print_file(const char* path)
{
    while (*path == ' ') path++;
    if (!g_file_io) {
        std::printf(CLR_RED "[ERROR] file_io 服务不可用\n" CLR_RESET);
        return;
    }
    char* content = nullptr;
    if (g_file_io->read(path, &content) == 0 && content) {
        std::printf("%s--- %s ---%s\n", CLR_DIM, path, CLR_RESET);
        std::printf("%s\n", content);
        std::printf(CLR_DIM "--- EOF ---\n" CLR_RESET);
        dstalk_free(content);
    } else {
        std::printf(CLR_RED "[ERROR] 无法读取: %s\n" CLR_RESET, path);
    }
}

static void list_files(const char* path)
{
    while (*path == ' ') path++;
    std::filesystem::path dir = *path ? std::filesystem::path(path) : std::filesystem::current_path();

    std::error_code ec;
    if (!std::filesystem::exists(dir, ec) || !std::filesystem::is_directory(dir, ec)) {
        std::printf(CLR_RED "[ERROR] 不是有效目录: %s\n" CLR_RESET, dir.string().c_str());
        return;
    }

    std::vector<std::filesystem::directory_entry> entries;
    for (const auto& entry : std::filesystem::directory_iterator(dir, ec)) {
        if (ec) break;
        entries.push_back(entry);
    }
    if (ec) {
        std::printf(CLR_RED "[ERROR] 无法列出目录: %s\n" CLR_RESET, dir.string().c_str());
        return;
    }

    std::sort(entries.begin(), entries.end(), [](const auto& a, const auto& b) {
        return a.path().filename().string() < b.path().filename().string();
    });

    std::printf("%s--- %s ---%s\n", CLR_DIM, dir.string().c_str(), CLR_RESET);
    for (const auto& entry : entries) {
        std::error_code status_ec;
        const bool is_dir = entry.is_directory(status_ec);
        std::printf("  %s%s%s\n", is_dir ? CLR_CYAN : "", entry.path().filename().string().c_str(), is_dir ? "/" CLR_RESET : "");
    }
}

static void handle_command(const char* line)
{
    if (!line || line[0] != '/') return;

    // /quit  —— 设置退出标志，让控制流自然回到 main 末尾
    if (std::strcmp(line, "/quit") == 0 || std::strcmp(line, "/q") == 0) {
        g_quit_requested = true;
        std::printf("再见!\n");
        return;
    }

    // /help
    if (std::strcmp(line, "/help") == 0 || std::strcmp(line, "/h") == 0) {
        print_help();
        return;
    }

    // /clear
    if (std::strcmp(line, "/clear") == 0) {
        if (g_session) g_session->clear();
        std::printf(CLR_GREEN "[OK] 会话已清空\n" CLR_RESET);
        return;
    }

    // /context
    if (std::strcmp(line, "/context") == 0) {
        if (g_session) {
            int count = 0;
            g_session->history(&count);
            int tokens = g_session->token_count();
            std::printf(CLR_DIM "消息条数: " CLR_RESET "%d  |  "
                        CLR_DIM "Token 估算: " CLR_RESET "%d\n",
                        count, tokens);
        }
        return;
    }

    // /status  —— 脱敏显示当前运行状态
    if (std::strcmp(line, "/status") == 0) {
        const char* provider = dstalk_config_get("ai.provider");
        if (!provider) provider = "ai.deepseek";
        const char* base_url = dstalk_config_get("api.base_url");
        if (!base_url) base_url = "https://api.deepseek.com/v1";
        const char* api_key  = dstalk_config_get("api.api_key");

        std::printf("  模型: %s\n", g_current_model.empty() ? "(未设置)" : g_current_model.c_str());
        std::printf("  base_url: %s\n", base_url ? base_url : "(未设置)");
        std::printf("  api_key: %s\n", (api_key && api_key[0]) ? "已设置" : "未设置");
        std::printf("  provider: %s\n", provider);
        std::printf("  AI 服务: %s\n", g_ai ? "就绪" : "不可用");
        std::printf("  Session 服务: %s\n", g_session ? "就绪" : "不可用");
        std::printf("  File IO 服务: %s\n", g_file_io ? "就绪" : "不可用");
        const dstalk_tools_service_t* tools = static_cast<const dstalk_tools_service_t*>(
            dstalk_service_query("tools", 1));
        std::printf("  Tools 服务: %s\n", tools ? "就绪" : "不可用");
        return;
    }

    // /model <name>
    if (std::strncmp(line, "/model ", 7) == 0) {
        const char* model = line + 7;
        while (*model == ' ') model++;
        if (*model == '\0') {
            std::printf(CLR_RED "[ERROR] /model 需要模型名\n" CLR_RESET);
            return;
        }
        if (g_ai) {
            g_ai->configure(nullptr, nullptr, nullptr, model, 0, 0.0);
            g_current_model = model;
            std::printf(CLR_GREEN "[OK] 模型已切换: %s\n" CLR_RESET, model);
        } else {
            std::printf(CLR_RED "[ERROR] AI 服务不可用\n" CLR_RESET);
        }
        return;
    }

    // /file list [path]
    if (std::strcmp(line, "/file list") == 0 || std::strncmp(line, "/file list ", 11) == 0) {
        const char* path = line + 10;
        list_files(path);
        return;
    }

    // /file show <path>
    if (std::strncmp(line, "/file show ", 11) == 0) {
        print_file(line + 11);
        return;
    }

    // /file read <path>
    if (std::strncmp(line, "/file read ", 11) == 0) {
        print_file(line + 11);
        return;
    }

    // /file write <path> <content...>
    if (std::strncmp(line, "/file write ", 12) == 0) {
        const char* rest = line + 12;
        while (*rest == ' ') rest++;
        const char* space = std::strchr(rest, ' ');
        if (!space) {
            std::printf(CLR_RED "[ERROR] 用法: /file write <path> <content>\n" CLR_RESET);
            return;
        }
        std::string path(rest, space - rest);
        const char* content = space + 1;
        while (*content == ' ') content++;
        if (g_file_io && g_file_io->write(path.c_str(), content) == 0) {
            std::printf(CLR_GREEN "[OK] 已写入: %s\n" CLR_RESET, path.c_str());
        } else {
            std::printf(CLR_RED "[ERROR] 写入失败: %s\n" CLR_RESET, path.c_str());
        }
        return;
    }

    // /save <path>
    if (std::strncmp(line, "/save ", 6) == 0) {
        const char* path = line + 6;
        while (*path == ' ') path++;
        if (g_session && g_session->save(path) == 0) {
            std::printf(CLR_GREEN "[OK] 会话已保存: %s\n" CLR_RESET, path);
        } else {
            std::printf(CLR_RED "[ERROR] 保存失败: %s\n" CLR_RESET, path);
        }
        return;
    }

    // /load <path>
    if (std::strncmp(line, "/load ", 6) == 0) {
        const char* path = line + 6;
        while (*path == ' ') path++;
        if (g_session && g_session->load(path) == 0) {
            std::printf(CLR_GREEN "[OK] 会话已恢复: %s\n" CLR_RESET, path);
        } else {
            std::printf(CLR_RED "[ERROR] 恢复失败: %s\n" CLR_RESET, path);
        }
        return;
    }

    std::printf(CLR_RED "未知命令: %s (输入 /help 查看帮助)\n" CLR_RESET, line);
}

// ---- 流式回调 ----
static int on_stream_token(const char* token, void* userdata)
{
    bool* first = static_cast<bool*>(userdata);
    if (*first) {
        std::printf(CLR_GREEN);
        *first = false;
    }
    std::printf("%s", token);
    std::fflush(stdout);
    return 0;
}

// ---- 主程序 ----
int main(int argc, char* argv[])
{
    // Windows: 启用 ANSI 转义码支持
#ifdef _WIN32
    HANDLE hOut = GetStdHandle(STD_OUTPUT_HANDLE);
    DWORD mode = 0;
    GetConsoleMode(hOut, &mode);
    SetConsoleMode(hOut, mode | ENABLE_VIRTUAL_TERMINAL_PROCESSING);
#endif

    // ---- C1: batch 模式检测 ----
    bool batch_mode = false;
    for (int i = 1; i < argc; ++i) {
        if (std::strcmp(argv[i], "--batch") == 0) {
            batch_mode = true;
            break;
        }
    }
#ifdef _WIN32
    if (!batch_mode && _isatty(_fileno(stdin)) == 0) batch_mode = true;
#else
    if (!batch_mode && isatty(fileno(stdin)) == 0) batch_mode = true;
#endif

    // ---- B1: 安装 Ctrl+C 处理 ----
#ifdef _WIN32
    SetConsoleCtrlHandler(on_console_event, TRUE);
#else
    signal(SIGINT, on_signal);
#endif

    // 查找配置文件
    const char* config_path = nullptr;
    if (argc >= 2) {
        // 跳过 --batch 标志
        for (int i = 1; i < argc; ++i) {
            if (std::strcmp(argv[i], "--batch") != 0) {
                config_path = argv[i];
                break;
            }
        }
    }
    if (!config_path) {
        const char* default_configs[] = {"config.toml", nullptr};
        for (int i = 0; default_configs[i]; i++) {
            FILE* f = nullptr;
#ifdef _WIN32
            fopen_s(&f, default_configs[i], "r");
#else
            f = fopen(default_configs[i], "r");
#endif
            if (f) {
                fclose(f);
                config_path = default_configs[i];
                break;
            }
        }
    }

    // 初始化主机（加载配置 + 自动扫描 plugins/ 目录加载插件）
    if (dstalk_init(config_path) != 0) {
        std::fprintf(stderr, CLR_RED "[dstalk] 初始化失败\n" CLR_RESET);
        return EXIT_INIT_FAIL;
    }

    // 查询插件服务
    const char* ai_provider = dstalk_config_get("ai.provider");
    if (!ai_provider) ai_provider = "ai.deepseek";
    g_ai      = static_cast<const dstalk_ai_service_t*>(dstalk_service_query(ai_provider, 1));
    g_session = static_cast<const dstalk_session_service_t*>(dstalk_service_query("session", 1));
    g_file_io = static_cast<const dstalk_file_io_service_t*>(dstalk_service_query("file_io", 1));

    if (!g_ai) {
        std::fprintf(stderr, CLR_RED "[dstalk] AI 服务未找到（请检查插件目录）\n" CLR_RESET);
    }
    if (!g_session) {
        std::fprintf(stderr, CLR_RED "[dstalk] Session 服务未找到\n" CLR_RESET);
    }

    // 自动从配置加载 AI 设置
    if (g_ai) {
        const char* base_url = dstalk_config_get("api.base_url");
        const char* api_key  = dstalk_config_get("api.api_key");
        const char* model    = dstalk_config_get("api.model");
        if (!base_url) base_url = "https://api.deepseek.com/v1";
        if (!model) model = "deepseek-v4-pro";
        g_ai->configure(ai_provider, base_url, api_key ? api_key : "", model, 4096, 0.7);
        g_current_model = model;   // A1: 记录当前模型名
    }

    if (!batch_mode) {
        std::printf("\n");
        print_banner();
        std::printf("\n");
    }

    char buffer[8192];
    while (true) {
        // B1: 检查退出标志
        if (g_quit_requested) break;

        // A1: 提示符带模型名（batch 模式不打印）
        if (!batch_mode) {
            std::printf(CLR_CYAN "[%s] " CLR_RESET CLR_YELLOW "> " CLR_RESET,
                        g_current_model.empty() ? "?" : g_current_model.c_str());
            std::fflush(stdout);
        }

        if (!std::fgets(buffer, sizeof(buffer), stdin)) break;

        // C3: fgets 截断检测
        if (!std::strchr(buffer, '\n') && !feof(stdin)) {
            std::fprintf(stderr, CLR_RED "[ERROR] 输入超过 8KB，已截断。建议用文件方式：dstalk --batch < file.txt\n" CLR_RESET);
            int c;
            while ((c = std::fgetc(stdin)) != '\n' && c != EOF) {}
        }

        // 去除末尾换行
        size_t len = std::strlen(buffer);
        while (len > 0 && (buffer[len-1] == '\n' || buffer[len-1] == '\r')) {
            buffer[--len] = '\0';
        }

        if (len == 0) continue;

        // 命令处理
        if (buffer[0] == '/') {
            handle_command(buffer);
            continue;
        }

        // AI 对话（通过插件服务 vtable）
        if (!g_ai || !g_session) {
            std::printf(CLR_RED "[ERROR] AI 或 Session 服务不可用\n" CLR_RESET);
            continue;
        }

        // 获取会话历史
        int history_count = 0;
        const dstalk_message_t* history = g_session->history(&history_count);

        bool first = true;
        dstalk_chat_result_t result = g_ai->chat_stream(
            history, history_count, buffer, on_stream_token, &first);

        if (result.ok) {
            std::printf(CLR_RESET "\n\n");
            // 将用户消息和 AI 回复添加到会话
            dstalk_message_t user_msg = {"user", buffer, nullptr, nullptr};
            g_session->add(&user_msg);
            dstalk_message_t ai_msg = {"assistant", result.content, nullptr, result.tool_calls_json};
            g_session->add(&ai_msg);
        } else {
            // A3: error 路径下需 NULL 保护；当前只取 result.error，content 未涉及
            std::printf(CLR_RESET "\n" CLR_RED "[ERROR] AI 调用失败: %s\n" CLR_RESET,
                        result.error ? result.error : "unknown error");
        }
        g_ai->free_result(&result);
    }

    // B2: 单一退出点，dstalk_shutdown 只在此调用
    dstalk_shutdown();
    return EXIT_OK;
}