include/ling3/chat_protocol.h

#pragma once

#include <nlohmann/json.hpp>
#include <algorithm>
#include <cmath>
#include <cctype>
#include <cstdint>
#include <functional>
#include <limits>
#include <stdexcept>
#include <string>
#include <string_view>
#include <utility>
#include <vector>

namespace ling3::chat {
using Json = nlohmann::json;
// Request bodies keep key order: the template renders tools as the client sent them.
using OrderedJson = nlohmann::ordered_json;
inline constexpr const char * kModel = "mindnano-ling3-tiny";
struct Error : std::runtime_error {
    int status;
    std::string code;
    Error(int s, std::string c, std::string m)
        : std::runtime_error(std::move(m)), status(s), code(std::move(c)) {}
};
inline Json ErrorBody(const Error & e) {
    return {{"error", {{"message", e.what()}, {"type", "invalid_request_error"},
                       {"param", nullptr}, {"code", e.code}}}};
}
struct Request {
    std::string prompt;
    std::size_t max_tokens = std::numeric_limits<std::size_t>::max();
    bool stream = false, include_usage = false;
    bool cache_prompt = true;
    bool reuse_generated_state = true;
    std::string cache_user;
    std::string session_id;
    double temperature = 0, top_p = 1;
    int top_k = 0;
    double repeat_penalty = 1;
    bool enable_thinking = false, flow_ack = false;
    std::uint32_t seed = 0;
    std::vector<std::string> stops;
    OrderedJson tools = OrderedJson::array();
    bool tool_calls = false, parallel_tool_calls = true; // parse calls from the output
};
inline std::size_t OutputBudget(std::size_t prompt, std::size_t capacity,
                                std::size_t requested) {
    if (!prompt || prompt > capacity)
        throw Error(400, "context_length_exceeded", "input exceeds allocated context of " +
                    std::to_string(capacity) + " tokens, or is empty");
    return std::min(requested, capacity - prompt);
}
// The template's tojson filter is Python json.dumps(ensure_ascii=False):
// ", " and ": " separators, keys in the order the client sent them.
inline std::string TemplateJson(const OrderedJson & v) {
    if (!v.is_object() && !v.is_array()) return v.dump();
    std::string s = v.is_object() ? "{" : "[";
    for (auto it = v.begin(); it != v.end(); ++it) {
        if (it != v.begin()) s += ", ";
        if (v.is_object()) s += OrderedJson(it.key()).dump() + ": ";
        s += TemplateJson(*it);
    }
    return s + (v.is_object() ? "}" : "]");
}
inline constexpr const char * kToolsHeader = "# Tools\n\nYou may call one or more functions to assist "
    "with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>";
inline constexpr const char * kToolsFooter = "\n</tools>\n\nIf none of the functions can be used, point it "
    "out. If the given question lacks the parameters required by the function, also point it out.\nIf you need "
    "to use a function, for each function call, output the function name and arguments within the following XML "
    "format:\n<tool_call>{function-name}\n<arg_key>{arg-key-1}</arg_key>\n<arg_value>{arg-value-1}</arg_value>\n"
    "<arg_key>{arg-key-2}</arg_key>\n<arg_value>{arg-value-2}</arg_value>\n...\n</tool_call>\n";
struct ToolCall { std::string name; OrderedJson arguments = OrderedJson::object(); };
// String values are raw text and others JSON, exactly as the template renders history.
inline std::string RenderToolCall(const ToolCall & call) {
    std::string s = "<tool_call>" + call.name + "\n";
    for (auto it = call.arguments.begin(); it != call.arguments.end(); ++it)
        s += "<arg_key>" + it.key() + "</arg_key>\n<arg_value>" +
             (it->is_string() ? it->get<std::string>() : TemplateJson(*it)) + "</arg_value>";
    return s + "\n</tool_call>";
}
inline Request Parse(const OrderedJson & j) {
    auto fail = [](std::string m) { throw Error(400, "invalid_parameter", std::move(m)); };
    if (!j.is_object()) fail("request must be a JSON object");
    for (auto it = j.begin(); it != j.end(); ++it) {
        const std::vector<std::string> supported {"model", "messages", "max_tokens",
            "max_completion_tokens", "stream", "stream_options", "temperature", "top_p",
            "seed", "stop", "n", "user", "cache_prompt", "top_k", "repeat_penalty",
            "enable_thinking", "flow_control", "session_id", "reuse_generated_state",
            "tools", "tool_choice", "parallel_tool_calls", "store"};
        if (std::find(supported.begin(), supported.end(), it.key()) == supported.end())
            fail("unsupported parameter: " + it.key());
    }
    if (!j.contains("model") || !j["model"].is_string() || j["model"] != kModel)
        throw Error(404, "model_not_found", std::string("model must be ") + kModel);
    Request r;
    if (j.contains("reuse_generated_state")) {
        if (!j["reuse_generated_state"].is_boolean()) fail("reuse_generated_state must be boolean");
        r.reuse_generated_state = j["reuse_generated_state"].get<bool>();
    }
    if (j.contains("session_id")) {
        if (!j["session_id"].is_string()) fail("session_id must be a string");
        r.session_id = j["session_id"].get<std::string>();
        if (r.session_id.empty() || r.session_id.size() > 256) fail("session_id must contain 1..256 bytes");
    }
    if (j.contains("top_k")) {
        if (!j["top_k"].is_number_integer() || j["top_k"] < 0 || j["top_k"] > 1000000)
            fail("top_k must be an integer in 0..1000000 (0 disables it)");
        r.top_k = j["top_k"].get<int>();
    }
    if (j.contains("repeat_penalty")) {
        if (!j["repeat_penalty"].is_number()) fail("repeat_penalty must be numeric");
        r.repeat_penalty = j["repeat_penalty"].get<double>();
        if (!std::isfinite(r.repeat_penalty) || r.repeat_penalty <= 0)
            fail("repeat_penalty must be finite and positive");
    }
    if (j.contains("enable_thinking")) {
        if (!j["enable_thinking"].is_boolean()) fail("enable_thinking must be boolean");
        r.enable_thinking = j["enable_thinking"].get<bool>();
    }
    if (j.contains("flow_control")) {
        if (j["flow_control"] != "none" && j["flow_control"] != "ack") fail("invalid flow_control");
        r.flow_ack = j["flow_control"] == "ack";
        if (r.flow_ack && j.value("stream", false) != true) fail("flow_control=ack requires streaming");
    }
    if (j.contains("cache_prompt")) {
        if (!j["cache_prompt"].is_boolean()) fail("cache_prompt must be boolean");
        r.cache_prompt = j["cache_prompt"].get<bool>();
    }
    if (j.contains("max_tokens") && j.contains("max_completion_tokens"))
        fail("set only one of max_tokens and max_completion_tokens");
    for (auto key : {"max_tokens", "max_completion_tokens"}) {
        if (!j.contains(key)) continue;
        const auto & value = j[key];
        if (!value.is_number_integer() ||
            (value.is_number_unsigned() ? value.get<std::uint64_t>() == 0
                                        : value.get<std::int64_t>() <= 0))
            fail("max_tokens/max_completion_tokens must be a positive integer");
        const auto count = value.get<std::uint64_t>();
        if (count > std::numeric_limits<std::size_t>::max()) fail("output budget exceeds addressable integer range");
        r.max_tokens = static_cast<std::size_t>(count);
    }
    if (j.contains("stream")) {
        if (!j["stream"].is_boolean()) fail("stream must be boolean");
        r.stream = j["stream"].get<bool>();
    }
    if (j.contains("stream_options") && !j["stream_options"].is_null()) {
        const auto & s = j["stream_options"];
        if (!r.stream || !s.is_object()) fail("stream_options requires stream=true");
        for (auto it = s.begin(); it != s.end(); ++it)
            if (it.key() != "include_usage" || !it.value().is_boolean())
                fail("only boolean stream_options.include_usage is supported");
        r.include_usage = s.value("include_usage", false);
    }
    for (auto key : {"temperature", "top_p"}) {
        if (!j.contains(key)) continue;
        if (!j[key].is_number()) fail(std::string(key) + " must be numeric");
        const auto x = j[key].get<double>();
        if (!std::isfinite(x) || x < 0 || x > (std::string(key) == "temperature" ? 2 : 1))
            fail(std::string(key) + " out of range");
        if (std::string(key) == "temperature") r.temperature = x;
        else { if (x == 0) fail("top_p must be > 0"); r.top_p = x; }
    }
    if (j.contains("seed")) {
        if (!j["seed"].is_number_integer() || j["seed"] < 0 || j["seed"] > UINT32_MAX)
            fail("seed must be a uint32 integer");
        r.seed = j["seed"].get<std::uint32_t>();
    }
    if (j.contains("n") && (!j["n"].is_number_integer() || j["n"] != 1)) fail("only n=1 is supported");
    // Agent clients send OpenAI's store flag; this server never stores completions.
    if (j.contains("store") && !j["store"].is_boolean()) fail("store must be boolean");
    if (j.contains("user") && !j["user"].is_string()) fail("user must be a string");
    if (j.contains("user")) r.cache_user = j["user"].get<std::string>();
    if (j.contains("stop") && !j["stop"].is_null()) {
        const auto s = j["stop"].is_string() ? OrderedJson::array({j["stop"]}) : j["stop"];
        if (!s.is_array() || s.size() > 4) fail("stop must be a string or at most 4 strings");
        for (const auto & v : s) {
            if (!v.is_string() || v.get_ref<const std::string &>().empty() ||
                v.get_ref<const std::string &>().size() > 256) fail("invalid stop string");
            r.stops.push_back(v.get<std::string>());
        }
    }
    if (j.contains("tools") && !j["tools"].is_null()) {
        if (!j["tools"].is_array()) fail("tools must be an array");
        for (const auto & t : j["tools"])
            if (!t.is_object() || !t.contains("type") || t["type"] != "function" || !t.contains("function") ||
                !t["function"].is_object() || !t["function"].contains("name") || !t["function"]["name"].is_string() ||
                t["function"]["name"].get_ref<const std::string &>().empty())
                fail("tools must be named function tools");
        r.tools = j["tools"];
    }
    // Calls are parsed after generation, so only auto and none can be honored.
    const bool no_tool_choice = !j.contains("tool_choice") || j["tool_choice"].is_null();
    if (!no_tool_choice && j["tool_choice"] != "auto" && j["tool_choice"] != "none")
        fail("tool_choice must be auto or none");
    if (j.contains("parallel_tool_calls")) {
        if (!j["parallel_tool_calls"].is_boolean()) fail("parallel_tool_calls must be boolean");
        r.parallel_tool_calls = j["parallel_tool_calls"].get<bool>();
    }
    r.tool_calls = !r.tools.empty() && (no_tool_choice || j["tool_choice"] == "auto");
    if (!j.contains("messages") || !j["messages"].is_array() || j["messages"].empty())
        fail("messages must be a non-empty array");
    struct Message { std::string role, content, reasoning; std::vector<ToolCall> calls; };
    std::vector<Message> messages;
    for (const auto & m : j["messages"]) {
        if (!m.is_object() || !m.contains("role") || !m["role"].is_string()) fail("messages require a string role");
        Message message {m["role"].get<std::string>()};
        const auto & role = message.role;
        if (role != "system" && role != "user" && role != "assistant" && role != "tool") fail("unsupported message role");
        std::vector<std::string> fields {"role", "content"};
        if (role == "assistant") fields.insert(fields.end(), {"tool_calls", "reasoning_content"});
        if (role == "tool") fields.insert(fields.end(), {"tool_call_id", "name"});
        for (auto it = m.begin(); it != m.end(); ++it)
            if (std::find(fields.begin(), fields.end(), it.key()) == fields.end()) fail("unsupported message field: " + it.key());
        // Only an assistant turn may omit content. Text parts are joined with
        // newlines, as vLLM flattens them for string-content templates.
        if (!m.contains("content") || m["content"].is_null()) {
            if (role != "assistant") fail("messages require string or text-part content");
        } else if (m["content"].is_string()) message.content = m["content"].get<std::string>();
        else if (m["content"].is_array()) {
            for (std::size_t k = 0; k < m["content"].size(); ++k) {
                const auto & part = m["content"][k];
                if (!part.is_object() || !part.contains("type") || part["type"] != "text" ||
                    !part.contains("text") || !part["text"].is_string()) fail("content parts must be text");
                message.content += (k ? "\n" : "") + part["text"].get<std::string>();
            }
        } else fail("messages require string or text-part content");
        if (role == "assistant" && m.contains("reasoning_content") && !m["reasoning_content"].is_null()) {
            if (!m["reasoning_content"].is_string()) fail("reasoning_content must be a string");
            message.reasoning = m["reasoning_content"].get<std::string>();
        }
        if (role == "assistant" && m.contains("tool_calls") && !m["tool_calls"].is_null()) {
            if (!m["tool_calls"].is_array()) fail("tool_calls must be an array");
            for (const auto & c : m["tool_calls"]) {
                if (!c.is_object() || !c.contains("function") || !c["function"].is_object() ||
                    !c["function"].contains("name") || !c["function"]["name"].is_string())
                    fail("tool_calls require a function name");
                ToolCall call {c["function"]["name"].get<std::string>()};
                // OpenAI clients send arguments as a JSON string; the template iterates a mapping.
                if (c["function"].contains("arguments")) {
                    const auto & a = c["function"]["arguments"];
                    call.arguments = a.is_string() ? OrderedJson::parse(a.get<std::string>(), nullptr, false) : a;
                    if (!call.arguments.is_object()) fail("tool call arguments must be a JSON object");
                }
                message.calls.push_back(std::move(call));
            }
        }
        messages.push_back(std::move(message));
    }
    if (messages.back().role != "user" && messages.back().role != "tool")
        fail("last message must be a user or tool message");
    // Bailing V3 template (chat_template.jinja). No hidden, server-owned conversation.
    const std::string thinking = r.enable_thinking ? "on" : "off";
    r.prompt = "<role>SYSTEM</role>";
    std::size_t start = 0;
    std::string system;
    bool system_mode = false;
    if (messages.front().role == "system") {
        system = messages.front().content;
        // The explicit API switch owns template control. Normalize conflicting
        // legacy instructions rather than producing two contradictory modes.
        for (const auto old : {std::string("detailed thinking off"), std::string("detailed thinking on")}) {
            std::size_t at = 0;
            while ((at = system.find(old, at)) != std::string::npos) {
                const auto replacement = "detailed thinking " + thinking;
                system.replace(at, old.size(), replacement); at += replacement.size();
            }
        }
        system_mode = system.find("detailed thinking " + thinking) != std::string::npos;
        start = 1;
    }
    if (!r.tools.empty()) {
        if (start) r.prompt += system + "\n";
        r.prompt += kToolsHeader;
        for (const auto & tool : r.tools) r.prompt += "\n" + TemplateJson(tool);
        r.prompt += kToolsFooter;
        if (!system_mode) r.prompt += "detailed thinking " + thinking;
    } else if (start) r.prompt += system + (system_mode ? "" : "\ndetailed thinking " + thinking);
    else r.prompt += "detailed thinking " + thinking;
    r.prompt += "<|role_end|>";
    auto strip = [](std::string s, bool left, bool right) {
        if (right) s.erase(s.find_last_not_of('\n') + 1);
        if (left) s.erase(0, std::min(s.find_first_not_of('\n'), s.size()));
        return s;
    };
    for (std::size_t i = start; i < messages.size(); ++i) {
        const auto & m = messages[i];
        if (m.role == "tool") {
            if (i == 0 || messages[i - 1].role != "tool") r.prompt += "<role>OBSERVATION</role>";
            r.prompt += "\n<tool_response>\n" + m.content + "\n</tool_response>";
            if (i + 1 == messages.size() || messages[i + 1].role != "tool") r.prompt += "<|role_end|>";
            continue;
        }
        if (m.role != "assistant") {
            r.prompt += (m.role == "user" ? "<role>HUMAN</role>" : "<role>SYSTEM</role>") + m.content + "<|role_end|>";
            continue;
        }
        auto content = m.content, reasoning = m.reasoning;
        if (reasoning.empty() && content.find("</think>") != std::string::npos) {
            reasoning = strip(content.substr(0, content.find("</think>")), false, true);
            if (const auto open = reasoning.rfind("<think>"); open != std::string::npos) reasoning.erase(0, open + 7);
            reasoning = strip(reasoning, true, false);
            content = strip(content.substr(content.rfind("</think>") + 8), true, false);
        }
        r.prompt += reasoning.empty() ? "<role>ASSISTANT</role>\n<think></think>" + content
            : "<role>ASSISTANT</role>\n<think>" + strip(reasoning, true, true) + "</think>" + strip(content, true, false);
        for (std::size_t k = 0; k < m.calls.size(); ++k)
            r.prompt += (k || !content.empty() ? "\n" : "") + RenderToolCall(m.calls[k]);
        r.prompt += "<|role_end|>";
    }
    r.prompt += r.enable_thinking ? "<role>ASSISTANT</role>\n<think>" : "<role>ASSISTANT</role>\n<think></think>";
    return r;
}

// Token pieces can split a UTF-8 character. Keep incomplete sequences until the
// next token; JSON encoders must never see partial bytes.
inline std::string TakeUtf8(std::string & pending, bool flush) {
    std::string out;
    std::size_t i = 0;
    while (i < pending.size()) {
        const auto c = static_cast<unsigned char>(pending[i]);
        std::size_t n = c < 0x80 ? 1 : c >= 0xc2 && c <= 0xdf ? 2 :
            c >= 0xe0 && c <= 0xef ? 3 : c >= 0xf0 && c <= 0xf4 ? 4 : 0;
        if (n && i + n > pending.size() && !flush) break;
        bool valid = n && i + n <= pending.size();
        for (std::size_t k = 1; valid && k < n; ++k)
            valid = (static_cast<unsigned char>(pending[i + k]) & 0xc0) == 0x80;
        if (valid && n >= 3) {
            const auto b = static_cast<unsigned char>(pending[i + 1]);
            valid = !((c == 0xe0 && b < 0xa0) || (c == 0xed && b >= 0xa0) ||
                      (c == 0xf0 && b < 0x90) || (c == 0xf4 && b >= 0x90));
        }
        if (valid) { out.append(pending, i, n); i += n; }
        else { out += "\xef\xbf\xbd"; ++i; }
    }
    pending.erase(0, i);
    return out;
}
class TextFilter {
    std::string pending_, utf8_;
    std::vector<std::string> stops_;
public:
    bool stopped = false;
    explicit TextFilter(std::vector<std::string> stops) : stops_(std::move(stops)) {}
    std::string Push(std::string_view piece, bool final = false) {
        if (stopped) return {};
        pending_ += piece;
        std::size_t end = std::string::npos;
        for (const auto & stop : stops_) end = std::min(end, pending_.find(stop));
        if (end != std::string::npos) { pending_.resize(end); stopped = true; final = true; }
        std::size_t keep = 0;
        if (!final) for (const auto & stop : stops_)
            for (std::size_t k = 1; k < stop.size() && k <= pending_.size(); ++k)
                if (pending_.compare(pending_.size() - k, k, stop, 0, k) == 0) keep = std::max(keep, k);
        utf8_ += pending_.substr(0, pending_.size() - keep);
        pending_.erase(0, pending_.size() - keep);
        return TakeUtf8(utf8_, final);
    }
};
// Declared string arguments stay raw text; other values parse as JSON and fall
// back to text, as llama.cpp's ling3 parser types them from the tool schema.
inline OrderedJson ArgumentValue(const OrderedJson & function, const std::string & key, const std::string & value) {
    bool text = false, other = false;
    std::function<void(const OrderedJson &)> scan = [&](const OrderedJson & schema) {
        if (!schema.is_object()) return;
        if (schema.contains("type"))
            for (const auto & type : schema["type"].is_array() ? schema["type"] : OrderedJson::array({schema["type"]}))
                (type == "string" ? text : other) = true;
        if (schema.contains("enum") && schema["enum"].is_array())
            for (const auto & choice : schema["enum"]) (choice.is_string() ? text : other) = true;
        for (const auto * union_key : {"anyOf", "oneOf"})
            if (schema.contains(union_key) && schema[union_key].is_array())
                for (const auto & branch : schema[union_key]) scan(branch);
    };
    if (function.contains("parameters") && function["parameters"].is_object() &&
        function["parameters"].contains("properties") && function["parameters"]["properties"].is_object() &&
        function["parameters"]["properties"].contains(key)) scan(function["parameters"]["properties"][key]);
    if (text && !other) return value;
    auto parsed = OrderedJson::parse(value, nullptr, false);
    return parsed.is_discarded() ? OrderedJson(value) : parsed;
}
// Bailing V3 calls: <tool_call>NAME\n<arg_key>K</arg_key>\n<arg_value>V</arg_value>...\n</tool_call>.
// Whitespace between the parts is optional and only declared tools are calls.
// Text after the last call is returned as trailing content, as llama.cpp does.
struct ToolCalls { bool ok = false; std::vector<ToolCall> calls; std::string trailing; };
inline ToolCalls ParseToolCalls(std::string_view text, const OrderedJson & tools, bool parallel) {
    ToolCalls result;
    auto space = [&] { while (!text.empty() && std::isspace(static_cast<unsigned char>(text.front()))) text.remove_prefix(1); };
    auto take = [&](std::string_view tag) { if (!text.starts_with(tag)) return false; text.remove_prefix(tag.size()); return true; };
    auto until = [&](std::string_view tag, std::string & out) {
        const auto at = text.find(tag);
        if (at == std::string_view::npos) return false;
        out = text.substr(0, at); text.remove_prefix(at + tag.size()); return true;
    };
    while (space(), take("<tool_call>")) {
        std::size_t n = 0;
        while (n < text.size() && !std::isspace(static_cast<unsigned char>(text[n])) && text[n] != '<') ++n;
        ToolCall call {std::string(text.substr(0, n))};
        text.remove_prefix(n);
        const OrderedJson * function = nullptr;
        for (const auto & tool : tools) if (tool["function"]["name"] == call.name) function = &tool["function"];
        if (!function) return {};
        while (space(), !take("</tool_call>")) {
            std::string key, value;
            if (!take("<arg_key>") || !until("</arg_key>", key)) return {};
            space();
            if (!take("<arg_value>") || !until("</arg_value>", value)) return {};
            call.arguments[key] = ArgumentValue(*function, key, value);
        }
        if (parallel || result.calls.empty()) result.calls.push_back(std::move(call));
    }
    if (result.calls.empty()) return {};
    result.ok = true;
    result.trailing = text;
    return result;
}
// Holds back generated text from the first <tool_call> so a call never streams
// as content; with thinking on this also ends the reasoning, as in llama.cpp.
// The newline the template writes before a call is held too and dropped.
class ToolCallFilter {
    std::string pending_, captured_;
    bool enabled_, newline_ = false;
public:
    explicit ToolCallFilter(bool enabled) : enabled_(enabled) {}
    const std::string & captured() const { return captured_; }
    // The exact held text, for output that does not parse as calls.
    std::string Released() const { return (newline_ ? "\n" : "") + captured_; }
    std::string Push(std::string_view piece, bool final = false) {
        if (!enabled_) return std::string(piece);
        if (!captured_.empty()) { captured_ += piece; return {}; }
        pending_ += piece;
        constexpr std::string_view tag = "<tool_call>";
        if (const auto at = pending_.find(tag); at != std::string::npos) {
            captured_ = pending_.substr(at);
            pending_.resize(at);
            newline_ = !pending_.empty() && pending_.back() == '\n';
            if (newline_) pending_.pop_back();
            return std::exchange(pending_, {});
        }
        std::size_t keep = 0;
        if (!final) {
            for (std::size_t k = std::min(tag.size() - 1, pending_.size()); k && !keep; --k)
                if (pending_.ends_with(tag.substr(0, k))) keep = k;
            if (keep < pending_.size() && pending_[pending_.size() - keep - 1] == '\n') ++keep;
        }
        auto out = pending_.substr(0, pending_.size() - keep);
        pending_.erase(0, pending_.size() - keep);
        return out;
    }
};
} // namespace ling3::chat