/* capability.c - 模型能力探测: 基线 -> 合并 -> 单独拆开 */
#include "capability.h"
#include "http.h"
#include "i18n.h"
#include "json.h"
#include "util.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/* 一张 1x1 的 PNG. 探测只需要"一个合法的图片字节流", 越小越好:
 * 服务端会把它缩放后按图计费, 尺寸再大也换不来更多信息.  */
static const char PROBE_PNG_B64[] = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1"
                                    "PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLv"
                                    "AAAAAElFTkSuQmCC";

/* 缓存: 按 base_url + model 索引. 同一个进程里换模型会自然重测,
 * 切回来则命中. 条数很少 (人手切的模型个数).  */
#define CAP_CACHE_MAX 8

typedef struct {
    char *base_url;
    char *model;
    Capabilities caps;
} CapEntry;

static CapEntry g_cache[CAP_CACHE_MAX];
static size_t g_cache_n;

void cap_cache_clear(void) {
    for (size_t i = 0; i < g_cache_n; i++) {
        free(g_cache[i].base_url);
        free(g_cache[i].model);
        g_cache[i].base_url = g_cache[i].model = NULL;
    }
    g_cache_n = 0;
}

const char *cap_name(int v) {
    switch (v) {
    case CAP_YES:
        return "支持";
    case CAP_NO:
        return "不支持";
    default:
        return "未测";
    }
}

/* 给人看的标签 (--verbose 的探测记录). cap_name() 是内部/测试用的中文常量,
 * 不随界面语言变, 所以这里另走一遍界面文案表.  */
static const char *cap_label(int v, int is_tools) {
    if (v == CAP_YES)
        return t("agent.caps.yes");
    if (v == CAP_NO)
        return t(is_tools ? "agent.caps.no.tools" : "agent.caps.no.vision");
    return t("agent.caps.unknown");
}

/* https://host/v1 + /chat/completions (base_url 末尾可能带斜杠) */
static char *chat_url(const char *base_url) {
    const char *bu = base_url ? base_url : "";
    size_t bl = strlen(bu);
    while (bl > 0 && bu[bl - 1] == '/')
        bl--;
    return xasprintf("%.*s/chat/completions", (int)bl, bu);
}

/* 一个试探请求的结论 */
/* 注意别叫 R_OK —— <unistd.h> 里那个是 access() 的权限位, 会撞名 */
typedef enum {
    PROBE_OK = 0,     /* 2xx: 接受 */
    PROBE_REJECT = 1, /* 400: 明确拒绝 (只有在基线通过时才有意义) */
    PROBE_UNKNOWN = 2 /* 网络、401/403/404/429/5xx……: 什么也说明不了 */
} ProbeResult;

static ProbeResult probe(const LlmConfig *cfg, const char *url,
                         const JsonValue *messages, int with_tools, int verbose,
                         long *out_status, char *errbuf, size_t errlen) {
    JsonValue *body = json_new_object();
    json_object_set_string(body, "model", cfg->model);
    json_object_set(body, "messages", json_clone(messages));
    /* 只让它回一个 token: 探测要的是"收不收", 不是回复内容.
     * 不设这个的话, 默认开着思考的模型会为了回答"1"生成一大段推理.  */
    json_object_set_number(body, "max_tokens", 1);
    json_object_set_bool(body, "stream", 0);
    if (with_tools) {
        JsonValue *tools = json_new_array();
        JsonValue *t = json_new_object();
        json_object_set_string(t, "type", "function");
        JsonValue *fn = json_new_object();
        json_object_set_string(fn, "name", "ligent_probe");
        json_object_set_string(fn, "description", "能力探测用, 永远不该被调用");
        JsonValue *params = json_new_object();
        json_object_set_string(params, "type", "object");
        json_object_set(fn, "parameters", params);
        json_object_set(t, "function", fn);
        json_array_push(tools, t);
        json_object_set(body, "tools", tools);
    }

    char *text = json_dump(body);
    json_free(body);

    char *resp_body = NULL;
    long status = 0;
    char local_err[256];
    local_err[0] = '\0';
    /* 探测不重试: 它只是一次判断, 服务端抖一下还有兜底 (一律按支持处理),
     * 不值得为此让每个进程启动都多等几个退避周期.  */
    HttpRetry no_retry = HTTP_RETRY_NONE;
    int rc =
        http_post(url, cfg->api_key, text, cfg->timeout, verbose, &no_retry,
                  &resp_body, &status, NULL, local_err, sizeof local_err);
    free(text);
    free(resp_body);
    if (out_status)
        *out_status = status;

    if (rc == 0)
        return PROBE_OK;
    if (status == 400)
        return PROBE_REJECT;

    if (errbuf && errlen && !errbuf[0] && local_err[0])
        snprintf(errbuf, errlen, "%s", local_err);
    return PROBE_UNKNOWN;
}

/* 探测用的消息: 一句 "1", 可选带一张图 */
static JsonValue *probe_messages(int with_image) {
    JsonValue *msgs = json_new_array();
    JsonValue *m = json_new_object();
    json_object_set_string(m, "role", "user");

    if (!with_image) {
        json_object_set_string(m, "content", "1");
        json_array_push(msgs, m);
        return msgs;
    }

    JsonValue *content = json_new_array();
    JsonValue *tp = json_new_object();
    json_object_set_string(tp, "type", "text");
    json_object_set_string(tp, "text", "1");
    json_array_push(content, tp);

    JsonValue *ip = json_new_object();
    json_object_set_string(ip, "type", "image_url");
    JsonValue *iu = json_new_object();
    char *data_url = xasprintf("data:image/png;base64,%s", PROBE_PNG_B64);
    json_object_set_string(iu, "url", data_url);
    free(data_url);
    json_object_set(ip, "image_url", iu);
    json_array_push(content, ip);

    json_object_set(m, "content", content);
    json_array_push(msgs, m);
    return msgs;
}

static void run_sequence(const LlmConfig *cfg, Capabilities *out, int verbose,
                         char *errbuf, size_t errlen) {
    out->vision = out->tools = CAP_UNKNOWN;
    out->probed = 1;
    out->conclusive = 0;

    char *url = chat_url(cfg->base_url);
    long status = 0;

    /* --- 1. 基线 --- */
    JsonValue *plain = probe_messages(0);
    ProbeResult base =
        probe(cfg, url, plain, 0, verbose, &status, errbuf, errlen);
    json_free(plain);

    if (base != PROBE_OK) {
        /* 连基线都不通, 后面的否定结论一概不作数 */
        if (verbose)
            fprintf(stderr, "%s",
                    tf("capability.note.baseline_failed", "status",
                       i18n_num(status), NULL));
        free(url);
        return;
    }
    out->conclusive = 1;

    /* --- 2. 合并试探: 图片 + tools --- */
    JsonValue *both = probe_messages(1);
    ProbeResult combined =
        probe(cfg, url, both, 1, verbose, &status, errbuf, errlen);
    json_free(both);

    if (combined == PROBE_OK) {
        /* 常见情况: 两样都支持, 两个请求收工 */
        out->vision = out->tools = CAP_YES;
        if (verbose)
            fprintf(stderr, "%s", t("capability.note.both_ok"));
        free(url);
        return;
    }

    /* --- 3. 拆开单独试: 才知道是哪一个不行 --- */
    JsonValue *img = probe_messages(1);
    ProbeResult vr = probe(cfg, url, img, 0, verbose, &status, errbuf, errlen);
    json_free(img);
    out->vision = (vr == PROBE_OK)
                      ? CAP_YES
                      : (vr == PROBE_REJECT ? CAP_NO : CAP_UNKNOWN);

    JsonValue *txt = probe_messages(0);
    ProbeResult tr = probe(cfg, url, txt, 1, verbose, &status, errbuf, errlen);
    json_free(txt);
    out->tools = (tr == PROBE_OK) ? CAP_YES
                                  : (tr == PROBE_REJECT ? CAP_NO : CAP_UNKNOWN);

    if (verbose)
        /* 用界面语言的标签, 不用 cap_name() —— 后者是给日志和断言用的中文常量
         */
        fprintf(stderr, "%s",
                tf("capability.note.split", "vision", cap_label(out->vision, 0),
                   "tools", cap_label(out->tools, 1), NULL));
    free(url);
}

int cap_cached(const LlmConfig *cfg, Capabilities *out) {
    const char *model = (cfg && cfg->model) ? cfg->model : "";
    const char *base = (cfg && cfg->base_url) ? cfg->base_url : "";
    for (size_t i = 0; i < g_cache_n; i++) {
        if (strcmp(g_cache[i].model, model) == 0 &&
            strcmp(g_cache[i].base_url, base) == 0) {
            if (out)
                *out = g_cache[i].caps;
            return 1;
        }
    }
    return 0;
}

int cap_get(const LlmConfig *cfg, int enable, Capabilities *out, int verbose,
            char *errbuf, size_t errlen) {
    if (!out)
        return 0;
    if (errbuf && errlen)
        errbuf[0] = '\0';
    if (!cfg) { /* 没有配置就什么都测不了, 按"未测"处理 */
        out->vision = out->tools = CAP_UNKNOWN;
        out->probed = 0;
        out->conclusive = 0;
        return 0;
    }

    const char *model = (cfg && cfg->model) ? cfg->model : "";
    const char *base = (cfg && cfg->base_url) ? cfg->base_url : "";

    if (!enable) {
        /* 探测被关掉: 不下任何结论, 一切都按"支持" */
        out->vision = out->tools = CAP_UNKNOWN;
        out->probed = 0;
        out->conclusive = 0;
        return 0;
    }

    for (size_t i = 0; i < g_cache_n; i++) {
        if (strcmp(g_cache[i].model, model) == 0 &&
            strcmp(g_cache[i].base_url, base) == 0) {
            *out = g_cache[i].caps;
            return 0;
        }
    }

    run_sequence(cfg, out, verbose, errbuf, errlen);

    /* 记进缓存; 满了就丢掉最老的一条 */
    if (g_cache_n == CAP_CACHE_MAX) {
        free(g_cache[0].base_url);
        free(g_cache[0].model);
        memmove(&g_cache[0], &g_cache[1],
                sizeof g_cache[0] * (CAP_CACHE_MAX - 1));
        g_cache_n--;
    }
    g_cache[g_cache_n].base_url = xstrdup(base);
    g_cache[g_cache_n].model = xstrdup(model);
    g_cache[g_cache_n].caps = *out;
    g_cache_n++;
    return 1;
}