/* llm.c - OpenAI 兼容接口: 请求构造、SSE 流式解析、tool_calls 增量拼接 */
#include "llm.h"
#include "config.h" /* THINKING_AUTO / _ON / _OFF */
#include "http.h"
#include "i18n.h"
#include "util.h"

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

/* ------------------------------------------------------------------ */
/* URL / 请求体                                                        */
/* ------------------------------------------------------------------ */

char *llm_endpoint_url(const char *base_url) {
    const char *base =
        (base_url && *base_url) ? base_url : "https://api.openai.com/v1";
    size_t n = strlen(base);
    while (n > 0 && base[n - 1] == '/')
        n--;
    /* 允许用户直接给完整的 completions 地址 */
    if (str_ends_with(base_url, "/chat/completions"))
        return xstrndup(base, n);
    return xasprintf("%.*s/chat/completions", (int)n, base);
}

JsonValue *llm_build_body(const LlmConfig *cfg, const JsonValue *messages,
                          const JsonValue *tools, int stream) {
    JsonValue *body = json_new_object();
    json_object_set_string(body, "model", cfg->model);
    json_object_set(body, "messages", json_clone(messages));
    if (cfg->temperature >= 0)
        json_object_set_number(body, "temperature", cfg->temperature);
    if (cfg->max_tokens > 0)
        json_object_set_number(body, "max_tokens", (double)cfg->max_tokens);
    /* 思考模式开关. AUTO 时整个字段都不发, 让服务端用自己的默认值
     *  (DeepSeek 默认是开启). 注意这些是请求体的顶层字段, 不是
     * OpenAI SDK 里的 extra_body 包装.  */
    if (cfg->thinking != THINKING_AUTO) {
        JsonValue *t = json_new_object();
        json_object_set_string(
            t, "type", cfg->thinking == THINKING_ON ? "enabled" : "disabled");
        json_object_set(body, "thinking", t);
    }
    if (cfg->reasoning_effort && *cfg->reasoning_effort)
        json_object_set_string(body, "reasoning_effort", cfg->reasoning_effort);
    if (tools && json_array_size(tools) > 0) {
        json_object_set(body, "tools", json_clone(tools));
        json_object_set_string(body, "tool_choice", "auto");
    }
    json_object_set_bool(body, "stream", stream ? 1 : 0);
    return body;
}

/* 单行 SSE 上限. 正常的 data: 行只有几十字节;
 * 出现超长"行"说明对端根本不是 SSE (例如忽略 stream 直接返回整包 JSON,
 * 而 JSON 里没有换行), 必须丢弃, 否则内存会被撑爆.  */
#define SSE_MAX_LINE (256u * 1024)
/* 单条回复累积内容上限 */
#define SSE_MAX_CONTENT (32u * 1024 * 1024)
/* 单个工具调用参数上限. 正常人写不出 8MB 的参数,
 * 超过基本可以断定是对端异常或恶意, 必须截住 */
#define SSE_MAX_TOOL_ARGS (8u * 1024 * 1024)

/* ------------------------------------------------------------------ */
/* 错误响应的可读化                                                    */
/* ------------------------------------------------------------------ */

/* 很多服务端把 message 写成 \uXXXX 转义, 直接打印很难看.
 * 这里尝试解析出 error.message, 失败则回退到原始文本.  */
static char *extract_error_message(const char *body) {
    if (!body || !*body)
        return NULL;
    JsonValue *root = json_parse(body, NULL, 0);
    if (!root)
        return NULL;
    char *msg = NULL;
    JsonValue *err = json_object_get(root, "error");
    if (err && err->type == JSON_OBJECT) {
        const char *m = json_get_string(err, "message");
        if (m && *m)
            msg = xstrdup(m);
    } else if (err && err->type == JSON_STRING) {
        msg = xstrdup(err->u.string);
    }
    if (!msg) {
        const char *m = json_get_string(root, "message");
        if (m && *m)
            msg = xstrdup(m);
    }
    json_free(root);
    return msg;
}

/* 401/403 时补一句提示. 服务端的原文经常把人引向错误的方向:
 * OpenRouter 对"格式不对的 key"回的是 "Missing Authentication header",
 * 听起来像是请求里没带认证头, 实际上带了, 只是 key 不属于这家.  */
static void append_auth_hint(char *errbuf, size_t errlen, const LlmConfig *cfg,
                             long status) {
    if (status != 401 && status != 403)
        return;

    const char *key = cfg->api_key ? cfg->api_key : "";
    size_t n = strlen(key);
    char masked[80];
    if (n == 0) {
        snprintf(masked, sizeof masked, "%s", t("llm.mask.empty"));
    } else {
        /* 只露前几个字符: 够认出是哪一把, 又不至于把密钥打进日志 */
        char head[8];
        snprintf(head, sizeof head, "%.*s", n <= 10 ? 4 : 7, key);
        snprintf(masked, sizeof masked, "%s",
                 tf("llm.mask.partial", "head", head, "n",
                    i18n_num((long long)n), NULL));
    }

    size_t cur = strlen(errbuf);
    if (cur + 4 >= errlen)
        return;
    snprintf(errbuf + cur, errlen - cur, "%s",
             tf("llm.err.auth_hint", "masked", masked, "url",
                cfg->base_url && *cfg->base_url ? cfg->base_url
                                                : t("llm.val.no_base_url"),
                NULL));
}

static void set_error_with_detail(char *errbuf, size_t errlen,
                                  const char *prefix, const char *body) {
    snprintf(errbuf, errlen, "%s", prefix);
    if (!body || !*body)
        return;
    size_t cur = strlen(errbuf);
    if (cur + 4 >= errlen)
        return;
    char *msg = extract_error_message(body);
    if (msg) {
        snprintf(errbuf + cur, errlen - cur, " — %s", msg);
        free(msg);
    } else {
        snprintf(errbuf + cur, errlen - cur, " — %.400s", body);
    }
}

/* ------------------------------------------------------------------ */
/* 流式累积器                                                          */
/* ------------------------------------------------------------------ */

typedef struct {
    char *id;
    char *name;
    StrBuf args;
    int started;
} ToolAcc;

typedef struct {
    StrBuf raw; /* 原始字节, 用于 "服务端不支持流式" 时的回退 */
    StrBuf content;
    StrBuf reasoning; /* reasoning_content 累积 (思考过程, 不显示) */
    ToolAcc *accs;
    size_t acc_count;
    size_t acc_cap;
    char *finish_reason;
    LlmDeltaFn on_delta;
    void *ud;
    LlmThinkFn on_think;
    void *think_ud;
    int got_event;
    int overflow; /* 内容超过上限 */
    int raw_truncated;
    char *error; /* 服务端返回的 error.message */
    JsonValue *usage;
    StrBuf linebuf;
} SseCtx;

static ToolAcc *acc_at(SseCtx *c, size_t index) {
    if (index >= c->acc_cap) {
        size_t cap = c->acc_cap ? c->acc_cap : 4;
        while (cap <= index)
            cap *= 2;
        c->accs = xrealloc(c->accs, cap * sizeof(ToolAcc));
        memset(c->accs + c->acc_cap, 0, (cap - c->acc_cap) * sizeof(ToolAcc));
        c->acc_cap = cap;
    }
    ToolAcc *a = &c->accs[index];
    if (!a->started) {
        a->started = 1;
        sb_init(&a->args);
        if (index + 1 > c->acc_count)
            c->acc_count = index + 1;
    }
    return a;
}

static void sse_ctx_free(SseCtx *c) {
    sb_free(&c->raw);
    sb_free(&c->content);
    sb_free(&c->reasoning);
    sb_free(&c->linebuf);
    for (size_t i = 0; i < c->acc_count; i++) {
        if (!c->accs[i].started)
            continue;
        free(c->accs[i].id);
        free(c->accs[i].name);
        sb_free(&c->accs[i].args);
    }
    free(c->accs);
    free(c->finish_reason);
    free(c->error);
    json_free(c->usage);
}

/* 处理一个 SSE data 负载 (不含 "data:" 前缀) */
static void handle_payload(SseCtx *c, const char *payload, size_t len) {
    if (len == 0)
        return;
    if (len == 6 && memcmp(payload, "[DONE]", 6) == 0)
        return;

    char *js = xstrndup(payload, len);
    char err[256];
    JsonValue *root = json_parse(js, err, sizeof err);
    free(js);
    if (!root)
        return;
    c->got_event = 1;

    JsonValue *errv = json_object_get(root, "error");
    if (errv && errv->type == JSON_OBJECT) {
        const char *m = json_get_string(errv, "message");
        if (m && !c->error)
            c->error = xstrdup(m);
        json_free(root);
        return;
    }
    JsonValue *usage = json_object_get(root, "usage");
    if (usage && usage->type == JSON_OBJECT && !c->usage)
        c->usage = json_clone(usage);

    JsonValue *choices = json_object_get(root, "choices");
    JsonValue *ch0 = json_array_get(choices, 0);
    if (!ch0) {
        json_free(root);
        return;
    }

    const char *fr = json_get_string(ch0, "finish_reason");
    if (fr && *fr && !c->finish_reason)
        c->finish_reason = xstrdup(fr);

    /* OpenAI 流式用 delta, 少数实现直接给 message */
    JsonValue *d = json_object_get(ch0, "delta");
    if (!d)
        d = json_object_get(ch0, "message");
    if (!d || d->type != JSON_OBJECT) {
        json_free(root);
        return;
    }

    const char *piece = json_get_string(d, "content");
    if (piece && *piece) {
        if (c->content.len + strlen(piece) > SSE_MAX_CONTENT) {
            c->overflow = 1;
        } else {
            sb_puts(&c->content, piece);
            if (c->on_delta)
                c->on_delta(c->ud, piece);
        }
    }

    /* 思考过程. DeepSeek 用 reasoning_content, 有些 OpenAI 兼容实现用
     * reasoning. 只累积、不回调 —— 用户明确要求不显示, 但它必须随
     * assistant 消息一起回传: 请求带 tools 时服务端要求这么做,
     * 漏了会在下一轮直接 400.  */
    const char *think = json_get_string(d, "reasoning_content");
    if (!think)
        think = json_get_string(d, "reasoning");
    if (think && *think) {
        if (c->reasoning.len + strlen(think) > SSE_MAX_CONTENT) {
            c->overflow = 1;
        } else {
            sb_puts(&c->reasoning, think);
            /* 不往正文回调里塞 —— 思维链不该被渲染. 单独走一路,
             * 让上层知道"模型开始推理了".  */
            if (c->on_think)
                c->on_think(c->think_ud, think);
        }
    }

    JsonValue *tcs = json_object_get(d, "tool_calls");
    if (tcs && tcs->type == JSON_ARRAY) {
        for (size_t i = 0; i < json_array_size(tcs); i++) {
            JsonValue *tc = json_array_get(tcs, i);
            if (!tc)
                continue;
            size_t idx = (size_t)json_get_number(tc, "index", (double)i);
            ToolAcc *a = acc_at(c, idx);
            const char *id = json_get_string(tc, "id");
            if (id && *id && !a->id)
                a->id = xstrdup(id);
            JsonValue *fn = json_object_get(tc, "function");
            if (fn && fn->type == JSON_OBJECT) {
                const char *name = json_get_string(fn, "name");
                if (name && *name) {
                    if (!a->name)
                        a->name = xstrdup(name);
                    else if (strcmp(a->name, name) != 0) {
                        /* 有些服务端会拆着发名字, 追加处理 */
                        char *joined = xasprintf("%s%s", a->name, name);
                        free(a->name);
                        a->name = joined;
                    }
                }
                const char *args = json_get_string(fn, "arguments");
                if (args) {
                    if (a->args.len + strlen(args) > SSE_MAX_TOOL_ARGS)
                        c->overflow = 1;
                    else
                        sb_puts(&a->args, args);
                }
            }
        }
    }
    json_free(root);
}

/* 按行切分, 逐条处理 "data:" */
static void feed(SseCtx *c, const char *data, size_t len) {
    sb_append(&c->linebuf, data, len);
    char *buf = c->linebuf.data;
    size_t start = 0;
    for (size_t i = 0; i < c->linebuf.len; i++) {
        if (buf[i] != '\n')
            continue;
        size_t line_len = i - start;
        if (line_len > 0 && buf[start + line_len - 1] == '\r')
            line_len--;
        const char *line = buf + start;
        if (line_len > 5 && memcmp(line, "data:", 5) == 0) {
            const char *p = line + 5;
            size_t n = line_len - 5;
            while (n > 0 && (*p == ' ' || *p == '\t')) {
                p++;
                n--;
            }
            handle_payload(c, p, n);
        }
        start = i + 1;
    }
    if (start > 0) {
        memmove(buf, buf + start, c->linebuf.len - start);
        c->linebuf.len -= start;
        buf[c->linebuf.len] = '\0';
    }

    /* 残行过长: 丢弃, 保证内存有界 */
    if (c->linebuf.len > SSE_MAX_LINE)
        sb_clear(&c->linebuf);
}

static size_t sse_write_cb(char *ptr, size_t size, size_t nmemb, void *ud) {
    SseCtx *c = ud;
    size_t n = size * nmemb;
    if (c->overflow)
        return 0; /* 超限, 中止传输 */

    if (c->raw.len < 65536) { /* 仅保留开头用于回退判断 */
        size_t take = n < 65536 - c->raw.len ? n : 65536 - c->raw.len;
        sb_append(&c->raw, ptr, take);
        if (take < n)
            c->raw_truncated = 1;
    } else {
        c->raw_truncated = 1;
    }
    feed(c, ptr, n);
    if (c->overflow)
        return 0;
    return n;
}

/* ------------------------------------------------------------------ */
/* 从整包 JSON 响应里抽出 assistant 消息                                */
/* ------------------------------------------------------------------ */

static LlmResponse *from_full_response(JsonValue *root, size_t raw_len) {
    (void)raw_len;
    JsonValue *errv = json_object_get(root, "error");
    if (errv && errv->type == JSON_OBJECT)
        return NULL;

    JsonValue *ch0 = json_array_get(json_object_get(root, "choices"), 0);
    if (!ch0)
        return NULL;
    JsonValue *msg = json_object_get(ch0, "message");
    if (!msg || msg->type != JSON_OBJECT)
        return NULL;

    LlmResponse *r = xcalloc(1, sizeof(LlmResponse));
    r->message = json_clone(msg);
    JsonValue *u = json_object_get(root, "usage");
    if (u && u->type == JSON_OBJECT)
        r->usage = json_clone(u);
    const char *fr = json_get_string(ch0, "finish_reason");
    if (fr)
        r->finish_reason = xstrdup(fr);
    return r;
}

/* 把累积结构整理成 LlmResponse */
static LlmResponse *from_sse(SseCtx *c) {
    LlmResponse *r = xcalloc(1, sizeof(LlmResponse));
    r->message = json_new_object();
    json_object_set_string(r->message, "role", "assistant");

    int has_tools = 0;
    for (size_t i = 0; i < c->acc_count; i++)
        if (c->accs[i].started) {
            has_tools = 1;
            break;
        }

    if (c->content.len > 0 || !has_tools)
        json_object_set_string(r->message, "content",
                               c->content.data ? c->content.data : "");
    else
        json_object_set_null(r->message, "content");

    /* 思考过程跟着 assistant 消息一起进历史. 用户看不到它, 但它必须原样回传:
     * 只要请求里带 tools, 服务端就要求把历史里的 reasoning_content 全带上,
     * 否则下一轮直接 400.  */
    if (c->reasoning.len > 0)
        json_object_set_string(r->message, "reasoning_content",
                               c->reasoning.data ? c->reasoning.data : "");

    if (has_tools) {
        JsonValue *arr = json_new_array();
        for (size_t i = 0; i < c->acc_count; i++) {
            ToolAcc *a = &c->accs[i];
            if (!a->started)
                continue;
            JsonValue *tc = json_new_object();
            if (a->id) {
                json_object_set_string(tc, "id", a->id);
            } else {
                char *gen = xasprintf("call_%zu", i);
                json_object_set_string(tc, "id", gen);
                free(gen);
            }
            json_object_set_string(tc, "type", "function");
            JsonValue *fn = json_new_object();
            json_object_set_string(fn, "name", a->name ? a->name : "");
            json_object_set_string(fn, "arguments",
                                   a->args.data ? a->args.data : "{}");
            json_object_set(tc, "function", fn);
            json_array_push(arr, tc);
        }
        json_object_set(r->message, "tool_calls", arr);
    }

    if (c->finish_reason)
        r->finish_reason = xstrdup(c->finish_reason);
    if (c->usage)
        r->usage = json_clone(c->usage);
    return r;
}

/* ------------------------------------------------------------------ */
/* 主入口                                                              */
/* ------------------------------------------------------------------ */

LlmResponse *llm_chat(const LlmConfig *cfg, const JsonValue *messages,
                      const JsonValue *tools, LlmDeltaFn on_delta, void *ud,
                      LlmThinkFn on_think, void *think_ud, char *errbuf,
                      size_t errlen) {
    static char errsink[1];
    if (!errbuf || !errlen) {
        errbuf = errsink;
        errlen = 1;
    }
    errbuf[0] = '\0';
    if (!cfg->api_key || !*cfg->api_key) {
        snprintf(errbuf, errlen, "%s", t("llm.err.no_api_key"));
        return NULL;
    }
    char *url = llm_endpoint_url(cfg->base_url);
    JsonValue *body = llm_build_body(cfg, messages, tools, cfg->stream);
    char *payload = json_dump(body);
    json_free(body);

    if (cfg->verbose)
        fprintf(stderr, "[ligent] POST %s\n[ligent] body: %.2000s\n", url,
                payload);

    LlmResponse *result = NULL;
    char http_err[512];
    long status = 0;

    if (cfg->stream) {
        SseCtx ctx;
        memset(&ctx, 0, sizeof ctx);
        sb_init(&ctx.raw);
        sb_init(&ctx.content);
        sb_init(&ctx.reasoning);
        sb_init(&ctx.linebuf);
        ctx.on_delta = on_delta;
        ctx.ud = ud;
        ctx.on_think = on_think;
        ctx.think_ud = think_ud;

        char *errbody = NULL;
        size_t delivered = 0;
        HttpRetry rt = HTTP_RETRY_DEFAULT;
        rt.max_retries = cfg->retries;
        int rc = http_post_stream(url, cfg->api_key, payload, cfg->timeout,
                                  cfg->verbose, &rt, sse_write_cb, &ctx,
                                  &errbody, &status, NULL, &delivered, http_err,
                                  sizeof http_err);
        if (g_ligent_interrupt) {
            result = xcalloc(1, sizeof(LlmResponse));
            result->interrupted = 1;
            result->message = json_new_object();
            json_object_set_string(result->message, "role", "assistant");
            json_object_set_string(result->message, "content",
                                   ctx.content.data ? ctx.content.data : "");
            if (ctx.reasoning.len > 0)
                json_object_set_string(result->message, "reasoning_content",
                                       ctx.reasoning.data ? ctx.reasoning.data
                                                          : "");
        } else if (ctx.overflow) {
            snprintf(errbuf, errlen, "%s",
                     tf("llm.err.overflow", "content",
                        i18n_num((long long)SSE_MAX_CONTENT), "args",
                        i18n_num((long long)SSE_MAX_TOOL_ARGS), NULL));
        } else if (rc != 0 && (ctx.content.len > 0 || ctx.got_event)) {
            /* 传输断了, 但已经收到并显示了部分内容.
             * 这些内容用户已经看到, 丢掉会让上下文和屏幕对不上;
             * 所以保留成一条"不完整"的消息, 交给上层提示.  */
            result = from_sse(&ctx);
            result->incomplete = 1;
            result->incomplete_reason = xstrdup(http_err);
        } else if (rc != 0) {
            if (errbody && *errbody) {
                set_error_with_detail(errbuf, errlen, http_err, errbody);
            } else if (ctx.error) {
                size_t cur = (size_t)snprintf(errbuf, errlen, "%s", http_err);
                if (cur + 4 < errlen)
                    snprintf(errbuf + cur, errlen - cur, " — %s", ctx.error);
            } else {
                snprintf(errbuf, errlen, "%s", http_err);
            }
            append_auth_hint(errbuf, errlen, cfg, status);
        } else if (ctx.error) {
            snprintf(errbuf, errlen, "%s",
                     tf("llm.err.server", "error", ctx.error, NULL));
        } else if (ctx.got_event) {
            result = from_sse(&ctx);
        } else {
            /* 服务端忽略了 stream 参数, 直接返回整包 JSON */
            char *raw = xstrndup(ctx.raw.data ? ctx.raw.data : "", ctx.raw.len);
            JsonValue *root = json_parse_ui(raw, http_err, sizeof http_err);
            if (root) {
                result = from_full_response(root, 0);
                json_free(root);
            }
            free(raw);
            if (!result && !errbuf[0]) {
                if (status >= 400)
                    snprintf(errbuf, errlen, "HTTP %ld", status);
                else
                    snprintf(errbuf, errlen, "%s",
                             t("llm.err.unparsable_response"));
            }
        }
        free(errbody);
        sse_ctx_free(&ctx);
    } else {
        char *resp = NULL;
        HttpRetry rt = HTTP_RETRY_DEFAULT;
        rt.max_retries = cfg->retries;
        int rc =
            http_post(url, cfg->api_key, payload, cfg->timeout, cfg->verbose,
                      &rt, &resp, &status, NULL, http_err, sizeof http_err);
        if (rc != 0) {
            if (resp && *resp)
                set_error_with_detail(errbuf, errlen, http_err, resp);
            else
                snprintf(errbuf, errlen, "%s", http_err);
            append_auth_hint(errbuf, errlen, cfg, status);
        } else {
            JsonValue *root =
                json_parse_ui(resp ? resp : "", http_err, sizeof http_err);
            if (root) {
                JsonValue *errv = json_object_get(root, "error");
                if (errv && errv->type == JSON_OBJECT) {
                    const char *m = json_get_string(errv, "message");
                    snprintf(errbuf, errlen, "%s",
                             tf("llm.err.server", "error",
                                m ? m : t("llm.val.unknown"), NULL));
                } else {
                    result = from_full_response(root, 0);
                    if (!result)
                        snprintf(errbuf, errlen, "%s", t("llm.err.no_choices"));
                }
                json_free(root);
            } else {
                char head[201];
                snprintf(head, sizeof head, "%.200s", http_err);
                snprintf(errbuf, errlen, "%s",
                         tf("llm.err.bad_json", "error", head, NULL));
            }
        }
        free(resp);
    }

    free(payload);
    free(url);
    return result;
}

void llm_response_free(LlmResponse *r) {
    if (!r)
        return;
    json_free(r->message);
    json_free(r->usage);
    free(r->finish_reason);
    free(r->incomplete_reason);
    free(r);
}