inferencelayer 0.2.13

//! The checkpoint's OWN Jinja `chat_template`, rendered in-process (feature `jinja`).
//!
//! The premise this module retires is recorded in [`super::chat`]: that HF chat templates need
//! Python-compatible filters "no Rust Jinja engine provides". They do not. `minijinja` plus a
//! handful of shims reproduces `transformers.apply_chat_template` **byte for byte** across the
//! served families — gated in `tests/chat_template_parity.rs` against a committed corpus of real
//! transformers renders (`bench/chat_template_fixtures/`).
//!
//! Byte-exactness is the whole point. The rendered prompt is tokenized, so one missing space after
//! a JSON comma re-segments the tool block into different ids and serves the model off the
//! distribution it was trained on. That shows up as "the model got a bit worse", never as a stack
//! trace — which is why the pieces below are matched to CPython's defaults rather than to Rust's.
//!
//! ## What transformers actually compiles (and what we match)
//!
//! `transformers.utils.chat_template_utils._cached_compile_jinja_template` builds an
//! `ImmutableSandboxedEnvironment(trim_blocks=True, lstrip_blocks=True, extensions=[…,
//! jinja2.ext.loopcontrols])` and installs three names: `tojson`, `raise_exception`,
//! `strftime_now`. Everything else is stock jinja2. So:
//!
//! - **`trim_blocks` / `lstrip_blocks`** — set here to the same values.
//! - **`keep_trailing_newline`** — left at the jinja2 default (`false`, i.e. the source's final
//!   newline is dropped at compile time). LiteRT-LM's port sets it `true`; doing that here would
//!   append a newline `apply_chat_template` never emits, and the corpus says so.
//! - **`tojson`** — `json.dumps` semantics: `", "` between elements, `": "` after keys, insertion
//!   order preserved, `ensure_ascii=False` by default. See [`tojson`].
//! - **`raise_exception`** — a recoverable [`TemplateError`], never a panic; the serving layer
//!   turns it into a 400.
//! - **`strftime_now`** — reads `now` from the render state so a date-stamping template
//!   (Llama-3, SmolLM3) is reproducible; falls back to the wall clock.
//! - **Python string/dict methods** (`.split()`, `.startswith()`, `.lstrip(chars)`, `.replace()`,
//!   `.get()`, …) via `minijinja-contrib`'s `pycompat` callback. Real templates call these as
//!   methods, not filters — Qwen3.5 and SmolLM3 do not render without it.
//! - **`lstrip` / `rstrip` filters** — char-set variants, ported from LiteRT-LM. jinja2 has no such
//!   filters, so these are additive; the method forms above are what templates in the wild use.
//!
//! One deliberate deviation from LiteRT-LM's port: the `none` **test** is left as minijinja's
//! builtin (`is_none`) rather than remapped to `is_undefined()`. jinja2's `x is none` is false for
//! an undefined `x`, and minijinja's builtin already agrees; remapping it would make
//! `{% if tools is none %}` — which every tool-aware template branches on — take the wrong arm.
//!
//! ## Ordering
//!
//! `serde_json/preserve_order` and minijinja's `preserve_order` are both enabled by the `jinja`
//! feature. Without them a tool schema's keys arrive alphabetized instead of in the order the
//! client sent, and `{{ tool | tojson }}` emits different bytes than Python's insertion-ordered
//! dicts. That is a silent prompt change, so the ordering is a correctness requirement, not a
//! cosmetic one.

use std::path::Path;

use minijinja::value::{Kwargs, Value as JValue, ValueKind};
use minijinja::{Environment, Error as JError, ErrorKind as JErrorKind, State};
use serde::Serialize;

/// What a template can express, probed from its source. The serving layer reads this to decide
/// whether the template already renders tools (so [`super::tools::tools_system_block`] is not
/// needed) and which content shape to hand it.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct TemplateCapabilities {
    /// The template renders a tool advertisement from a `tools` variable.
    pub supports_tools: bool,
    /// It replays an assistant turn's `tool_calls`.
    pub supports_tool_calls: bool,
    /// It has a `system` arm (as opposed to folding system text into the first user turn).
    pub supports_system_role: bool,
    /// It loops over `tool_calls`, so more than one call per assistant turn round-trips.
    pub supports_parallel_tool_calls: bool,
    /// It reads `tool_call_id` (the OpenAI/Llama correlation field) off tool results.
    pub supports_tool_call_id: bool,
    /// `content` must be a list of typed blocks (`[{"type":"text","text":…}]`), not a string.
    /// Detected BEHAVIORALLY — render once each way and see which round-trips — so it holds across
    /// families instead of tracking a regex against whatever this month's VLM template looks like.
    pub requires_typed_content: bool,
    /// The template honours `is_appending_to_prefill`: the caller may append only the new turn
    /// instead of re-rendering the whole history. The consumer is the session object (plan item
    /// S1); recorded now because it is what makes prefix reuse pay on multi-turn.
    pub supports_single_turn: bool,
}

/// A template failure: a compile error, or a render error (including `raise_exception`).
#[derive(Debug, Clone)]
pub struct TemplateError {
    pub message: String,
    pub stage: TemplateStage,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TemplateStage {
    /// The source did not parse — the checkpoint ships a template we cannot compile.
    Compile,
    /// Rendering failed: `raise_exception`, an unsupported construct, or bad inputs.
    Render,
}

impl std::fmt::Display for TemplateError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let what = match self.stage {
            TemplateStage::Compile => "chat template compile",
            TemplateStage::Render => "chat template render",
        };
        write!(f, "{what}: {}", self.message)
    }
}

impl std::error::Error for TemplateError {}

impl TemplateError {
    fn from_jinja(stage: TemplateStage, e: &JError) -> Self {
        // minijinja's Display is one line; the chain carries the `raise_exception` message on
        // templates that wrap it, so walk it and keep everything the template author wrote.
        let mut message = e.to_string();
        let mut src: Option<&(dyn std::error::Error + 'static)> = std::error::Error::source(e);
        while let Some(e) = src {
            message.push_str(": ");
            message.push_str(&e.to_string());
            src = std::error::Error::source(e);
        }
        Self { message, stage }
    }
}

/// Everything `apply_chat_template` puts in the render context.
///
/// The field set mirrors transformers' own call — `messages`, `tools`, `documents`,
/// `add_generation_prompt`, then the tokenizer's `special_tokens_map` — because a template that
/// branches on `tools is none` behaves differently when the variable is absent versus null.
#[derive(Debug, Clone, Default)]
pub struct RenderInputs {
    /// The conversation, as a JSON array of message objects.
    pub messages: serde_json::Value,
    /// Tool schemas. `None` renders as JSON `null` (what `apply_chat_template(tools=None)` does),
    /// NOT as an undefined variable.
    pub tools: Option<serde_json::Value>,
    /// RAG documents, same null-vs-undefined rule as `tools`.
    pub documents: Option<serde_json::Value>,
    pub add_generation_prompt: bool,
    /// Unix seconds for [`strftime_now`]. `None` falls back to the wall clock; supply it whenever
    /// the render must be reproducible.
    pub now: Option<i64>,
    /// Extra context: `bos_token`, `eos_token`, `enable_thinking`, `date_string`, … Merged last,
    /// so it wins over the fields above.
    pub extra: serde_json::Map<String, serde_json::Value>,
}

/// A compiled chat template plus the capabilities probed from its source.
#[derive(Clone)]
pub struct ChatTemplate {
    env: Environment<'static>,
    source: String,
    caps: TemplateCapabilities,
}

const TEMPLATE_NAME: &str = "chat";

impl ChatTemplate {
    /// Compile `source`. A parse failure is returned, not panicked: a checkpoint shipping a
    /// template we cannot compile must fall back to the arch renderer, not take the server down.
    pub fn new(source: impl Into<String>) -> Result<Self, TemplateError> {
        let source = source.into();
        let mut env = create_env();
        env.add_template_owned(TEMPLATE_NAME, rewrite_generation_blocks(&source))
            .map_err(|e| TemplateError::from_jinja(TemplateStage::Compile, &e))?;
        let caps = detect_capabilities(&source);
        Ok(Self { env, source, caps })
    }

    pub fn source(&self) -> &str {
        &self.source
    }

    pub fn capabilities(&self) -> TemplateCapabilities {
        self.caps
    }

    /// Render the conversation. Mirrors `apply_chat_template`'s context construction exactly; see
    /// [`RenderInputs`].
    pub fn render(&self, inputs: &RenderInputs) -> Result<String, TemplateError> {
        let tmpl = self
            .env
            .get_template(TEMPLATE_NAME)
            .map_err(|e| TemplateError::from_jinja(TemplateStage::Compile, &e))?;

        let mut ctx = serde_json::Map::new();
        ctx.insert(
            "messages".into(),
            match &inputs.messages {
                serde_json::Value::Null => serde_json::Value::Array(Vec::new()),
                other => other.clone(),
            },
        );
        // `tools`/`documents` are ALWAYS defined (null when absent) — that is what transformers
        // passes, and `{% if tools is none %}` is a real branch in most tool-aware templates.
        ctx.insert(
            "tools".into(),
            inputs.tools.clone().unwrap_or(serde_json::Value::Null),
        );
        ctx.insert(
            "documents".into(),
            inputs.documents.clone().unwrap_or(serde_json::Value::Null),
        );
        ctx.insert(
            "add_generation_prompt".into(),
            serde_json::Value::Bool(inputs.add_generation_prompt),
        );
        if let Some(now) = inputs.now {
            ctx.insert("now".into(), serde_json::Value::from(now));
        }
        for (k, v) in &inputs.extra {
            ctx.insert(k.clone(), v.clone());
        }

        tmpl.render(JValue::from_serialize(serde_json::Value::Object(ctx)))
            .map_err(|e| TemplateError::from_jinja(TemplateStage::Render, &e))
    }
}

/// A checkpoint's chat template together with the special tokens the render context needs.
pub struct CheckpointChat {
    pub template: ChatTemplate,
    /// HF's `special_tokens_map` — `bos_token`, `eos_token`, … normalized to plain strings. These
    /// are exactly the names `apply_chat_template` injects, and templates like gemma-4's open with
    /// `{{ bos_token }}`.
    pub special_tokens: serde_json::Map<String, serde_json::Value>,
}

/// The special-token names transformers exposes to a template.
const SPECIAL_TOKEN_KEYS: &[&str] = &[
    "bos_token",
    "eos_token",
    "unk_token",
    "sep_token",
    "pad_token",
    "cls_token",
    "mask_token",
    "additional_special_tokens",
];

/// Load `dir`'s chat template, if it ships one.
///
/// Both on-disk shapes are handled: the current standalone `chat_template.jinja`, and the legacy
/// `chat_template` key inside `tokenizer_config.json` (either a bare string or the
/// `[{"name":…,"template":…}]` multi-template list, from which `default` is taken). `Ok(None)`
/// means "this checkpoint has no template" — the caller falls back to
/// [`super::chat::render_chat_prompt`].
pub fn load_checkpoint_chat(dir: &Path) -> Result<Option<CheckpointChat>, TemplateError> {
    let tok_cfg: Option<serde_json::Value> =
        std::fs::read_to_string(dir.join("tokenizer_config.json"))
            .ok()
            .and_then(|s| serde_json::from_str(&s).ok());

    let source = match std::fs::read_to_string(dir.join("chat_template.jinja")) {
        Ok(s) => Some(s),
        Err(_) => tok_cfg
            .as_ref()
            .and_then(|c| c.get("chat_template"))
            .and_then(template_from_config_value),
    };
    let Some(source) = source else {
        return Ok(None);
    };

    let mut special_tokens = serde_json::Map::new();
    // `special_tokens_map.json` is authoritative when present; `tokenizer_config.json` carries the
    // same keys on checkpoints that ship only the one file (gemma-4, Qwen3.5).
    for file in ["tokenizer_config.json", "special_tokens_map.json"] {
        let Ok(raw) = std::fs::read_to_string(dir.join(file)) else {
            continue;
        };
        let Ok(v) = serde_json::from_str::<serde_json::Value>(&raw) else {
            continue;
        };
        for key in SPECIAL_TOKEN_KEYS {
            if let Some(tok) = v.get(*key)
                && let Some(normalized) = normalize_special_token(tok)
            {
                special_tokens.insert((*key).to_string(), normalized);
            }
        }
    }

    Ok(Some(CheckpointChat {
        template: ChatTemplate::new(source)?,
        special_tokens,
    }))
}

/// `chat_template` in `tokenizer_config.json` is either the source string or the legacy
/// `[{"name": "default", "template": …}]` list.
fn template_from_config_value(v: &serde_json::Value) -> Option<String> {
    match v {
        serde_json::Value::String(s) => Some(s.clone()),
        serde_json::Value::Array(entries) => entries
            .iter()
            .find(|e| e.get("name").and_then(|n| n.as_str()) == Some("default"))
            .or_else(|| entries.first())
            .and_then(|e| e.get("template"))
            .and_then(|t| t.as_str())
            .map(str::to_string),
        _ => None,
    }
}

/// HF writes a special token either as a plain string or as a serialized `AddedToken`
/// (`{"content": "<eos>", "lstrip": false, …}`). Templates interpolate the STRING, so unwrap it.
fn normalize_special_token(v: &serde_json::Value) -> Option<serde_json::Value> {
    match v {
        serde_json::Value::String(_) => Some(v.clone()),
        serde_json::Value::Object(o) => o.get("content").filter(|c| c.is_string()).cloned(),
        serde_json::Value::Array(items) => Some(serde_json::Value::Array(
            items.iter().filter_map(normalize_special_token).collect(),
        )),
        _ => None,
    }
}

// ============================================================================
// Capability probe
// ============================================================================

/// Probe what `source` supports. The `requires_typed_content` half is behavioral (render the same
/// message twice, once with string content and once with typed blocks, and see which one survives);
/// the rest are source scans, which is all the information a Jinja source carries about them.
fn detect_capabilities(source: &str) -> TemplateCapabilities {
    let mut caps = TemplateCapabilities::default();
    let env = create_env();
    let compilable = rewrite_generation_blocks(source);

    if let Ok(tmpl) = env.template_from_str(&compilable) {
        if tmpl.undeclared_variables(true).contains("tools") {
            caps.supports_tools = true;
        }

        const PROBE: &str = "test content";
        let try_render = |content: serde_json::Value| -> bool {
            let ctx = serde_json::json!({
                "messages": [{"role": "user", "content": content}],
                "add_generation_prompt": false,
                "tools": [],
                "documents": null,
            });
            tmpl.render(JValue::from_serialize(&ctx))
                .map(|s| s.contains(PROBE))
                .unwrap_or(false)
        };
        let string_works = try_render(serde_json::json!(PROBE));
        let typed_works = try_render(serde_json::json!([{"type": "text", "text": PROBE}]));
        if !string_works && typed_works {
            caps.requires_typed_content = true;
        }
    }

    if source.contains("tool_calls") {
        caps.supports_tool_calls = true;
    }
    if source.contains("tool_call_id") {
        caps.supports_tool_call_id = true;
    }
    if !caps.supports_tools && source.contains("tools") {
        caps.supports_tools = true;
    }
    if source.contains("system") {
        caps.supports_system_role = true;
    }
    if caps.supports_tool_calls && source.contains("for") {
        caps.supports_parallel_tool_calls = true;
    }
    if source.contains("is_appending_to_prefill") {
        caps.supports_single_turn = true;
    }
    caps
}

// ============================================================================
// The environment and its shims
// ============================================================================

// ============================================================================
// CPython stringification
// ============================================================================

/// Python's `str(x)` — what jinja2 writes for `{{ x }}` and what the `string` filter (`soft_str`)
/// returns.
///
/// This is not a cosmetic difference. SmolLM3's tool advertisement is `{{ tool | string }}` on a
/// dict, so the block the model was trained on is a **Python dict repr** — single quotes, `True`,
/// `None` — not JSON. Rendering minijinja's JSON-ish default there changes every tool prompt.
/// A top-level string stringifies to itself (`str("a") == "a"`); only nested strings get quoted,
/// exactly as in Python.
fn python_str(v: &JValue) -> String {
    if v.is_undefined() {
        return String::new(); // jinja2's Undefined.__str__
    }
    if v.is_none() {
        return "None".into();
    }
    match v.kind() {
        ValueKind::String => v.as_str().unwrap_or_default().to_owned(),
        ValueKind::Bool => bool_repr(v).into(),
        ValueKind::Seq | ValueKind::Iterable | ValueKind::Map => python_repr(v),
        _ => v.to_string(),
    }
}

/// Python's `repr(x)`, used for values NESTED inside a container.
fn python_repr(v: &JValue) -> String {
    if v.is_none() {
        return "None".into();
    }
    match v.kind() {
        ValueKind::String => python_quote(v.as_str().unwrap_or_default()),
        ValueKind::Bool => bool_repr(v).into(),
        ValueKind::Seq | ValueKind::Iterable => {
            let items = v
                .try_iter()
                .map(|it| it.map(|x| python_repr(&x)).collect::<Vec<_>>())
                .unwrap_or_default();
            format!("[{}]", items.join(", "))
        }
        ValueKind::Map => {
            let entries = v
                .try_iter()
                .map(|it| {
                    it.map(|k| {
                        let val = v.get_item(&k).unwrap_or(JValue::UNDEFINED);
                        format!("{}: {}", python_repr(&k), python_repr(&val))
                    })
                    .collect::<Vec<_>>()
                })
                .unwrap_or_default();
            format!("{{{}}}", entries.join(", "))
        }
        _ => v.to_string(),
    }
}

fn bool_repr(v: &JValue) -> &'static str {
    if v.is_true() { "True" } else { "False" }
}

/// CPython's `repr()` quoting for `str`: single quotes unless that would need escaping and double
/// quotes would not; backslash, the quote char and the C escapes escaped; other control characters
/// as `\xNN`. Non-ASCII is left alone (Python 3 reprs are unicode).
fn python_quote(s: &str) -> String {
    let quote = if s.contains('\'') && !s.contains('"') {
        '"'
    } else {
        '\''
    };
    let mut out = String::with_capacity(s.len() + 2);
    out.push(quote);
    for c in s.chars() {
        match c {
            '\\' => out.push_str("\\\\"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            c if c == quote => {
                out.push('\\');
                out.push(c);
            }
            c if (c as u32) < 0x20 || c as u32 == 0x7f => {
                out.push_str(&format!("\\x{:02x}", c as u32));
            }
            c => out.push(c),
        }
    }
    out.push(quote);
    out
}

/// Rewrite transformers' `{% generation %}` … `{% endgeneration %}` into `{% if true %}` …
/// `{% endif %}`.
///
/// That tag comes from the `AssistantTracker` jinja2 extension transformers compiles in. Its ONLY
/// effect is to record where the assistant span starts and ends so `return_assistant_tokens_mask`
/// can build a mask — the body renders verbatim either way, and we do not serve that flag. minijinja
/// has no custom-statement hook, so the tag is rewritten to a block statement with identical
/// whitespace semantics (both are block tags, so `trim_blocks`/`lstrip_blocks` apply the same) and
/// a body that always renders. Whitespace-control markers (`{%-`, `-%}`) are carried across.
///
/// SmolLM3 ships this tag; without the rewrite its template does not compile at all.
fn rewrite_generation_blocks(source: &str) -> String {
    let mut out = String::with_capacity(source.len());
    let mut i = 0;
    while i < source.len() {
        let Some(rel) = source[i..].find("{%") else {
            out.push_str(&source[i..]);
            return out;
        };
        let start = i + rel;
        out.push_str(&source[i..start]);
        let Some(rel_end) = source[start..].find("%}") else {
            out.push_str(&source[start..]);
            return out;
        };
        let end = start + rel_end + 2;
        let inner = &source[start + 2..end - 2];
        let (lead_dash, rest) = match inner.strip_prefix('-') {
            Some(r) => (true, r),
            None => (false, inner),
        };
        let (trail_dash, rest) = match rest.strip_suffix('-') {
            Some(r) => (true, r),
            None => (false, rest),
        };
        match rest.trim() {
            "generation" => push_stmt(&mut out, lead_dash, "if true", trail_dash),
            "endgeneration" => push_stmt(&mut out, lead_dash, "endif", trail_dash),
            _ => out.push_str(&source[start..end]),
        }
        i = end;
    }
    out
}

fn push_stmt(out: &mut String, lead_dash: bool, stmt: &str, trail_dash: bool) {
    out.push_str(if lead_dash { "{%- " } else { "{% " });
    out.push_str(stmt);
    out.push_str(if trail_dash { " -%}" } else { " %}" });
}

fn create_env() -> Environment<'static> {
    let mut env = Environment::new();
    // Matches `ImmutableSandboxedEnvironment(trim_blocks=True, lstrip_blocks=True)`.
    // `keep_trailing_newline` stays at the jinja2 default (false) — see the module docs.
    env.set_trim_blocks(true);
    env.set_lstrip_blocks(true);
    // Python string/dict methods (`.split()`, `.startswith()`, `.get()`, …). Templates call these
    // as methods; without the callback Qwen3.5 and SmolLM3 do not render at all.
    env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback);
    // `{{ x }}` must stringify the way CPython does — `True`, `None`, and single-quoted dict
    // reprs — because that is what jinja2 wrote into the prompts these models were trained on.
    env.set_formatter(|out, state, value| {
        match value.kind() {
            ValueKind::Bool
            | ValueKind::None
            | ValueKind::Map
            | ValueKind::Seq
            | ValueKind::Iterable => {
                minijinja::escape_formatter(out, state, &JValue::from(python_str(value)))
            }
            // Strings and numbers already render identically; keep the default path so auto-escape
            // and the fast paths stay untouched.
            _ => minijinja::escape_formatter(out, state, value),
        }
    });
    // jinja2's `string` filter is `soft_str`, i.e. `str()`. SmolLM3 advertises tools through it.
    env.add_filter("string", |v: JValue| python_str(&v));
    env.add_function("strftime_now", strftime_now);
    env.add_function("raise_exception", raise_exception);
    env.add_filter("tojson", tojson);
    env.add_filter("lstrip", lstrip);
    env.add_filter("rstrip", rstrip);
    env
}

/// `strftime_now(fmt)` — format `now` (unix seconds, from the render context) with a C `strftime`
/// spec. Absent `now`, use the wall clock.
///
/// UTC, not local time: the reproducible path is the one that matters (the fixture corpus pins a
/// timestamp and asserts the bytes), and a server's local zone must not change the prompt.
fn strftime_now(state: &State, format: String) -> Result<String, JError> {
    let now = state.lookup("now");
    let timestamp = match now {
        Some(v) if !v.is_undefined() && !v.is_none() => i64::try_from(v)
            .map_err(|_| JError::new(JErrorKind::InvalidOperation, "`now` must be unix seconds"))?,
        _ => wall_clock_seconds()?,
    };
    let dt = chrono::DateTime::from_timestamp(timestamp, 0)
        .ok_or_else(|| JError::new(JErrorKind::InvalidOperation, "timestamp out of range"))?;
    Ok(dt.format(&format).to_string())
}

#[cfg(not(target_arch = "wasm32"))]
fn wall_clock_seconds() -> Result<i64, JError> {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs() as i64)
        .map_err(|e| JError::new(JErrorKind::InvalidOperation, e.to_string()))
}

/// wasm32 has no wall clock in `std::time`; a browser render must supply `now` explicitly rather
/// than trap.
#[cfg(target_arch = "wasm32")]
fn wall_clock_seconds() -> Result<i64, JError> {
    Err(JError::new(
        JErrorKind::InvalidOperation,
        "strftime_now needs an explicit `now` on this target (no wall clock)",
    ))
}

/// `raise_exception(msg)` — how HF templates reject an unsupported role order. Surfaces as a
/// render error the handler turns into a 400.
fn raise_exception(msg: String) -> Result<String, JError> {
    Err(JError::new(JErrorKind::InvalidOperation, msg))
}

/// A `serde_json` formatter reproducing CPython's `json.dumps` default separators: `", "` between
/// elements and object entries, `": "` after a key. serde's compact default omits both, which
/// changes the prompt bytes and therefore the token ids.
struct SpaceFormatter;

impl serde_json::ser::Formatter for SpaceFormatter {
    fn begin_array_value<W>(&mut self, writer: &mut W, first: bool) -> std::io::Result<()>
    where
        W: ?Sized + std::io::Write,
    {
        if !first {
            writer.write_all(b", ")?;
        }
        Ok(())
    }

    fn begin_object_key<W>(&mut self, writer: &mut W, first: bool) -> std::io::Result<()>
    where
        W: ?Sized + std::io::Write,
    {
        if !first {
            writer.write_all(b", ")?;
        }
        Ok(())
    }

    fn begin_object_value<W>(&mut self, writer: &mut W) -> std::io::Result<()>
    where
        W: ?Sized + std::io::Write,
    {
        writer.write_all(b": ")
    }
}

/// `value | tojson(indent=…, ensure_ascii=…)` — transformers' override of jinja2's builtin, which
/// is plain `json.dumps` (no HTML escaping, `ensure_ascii=False` by default, `indent=None`).
///
/// An unrecognized keyword is an error rather than a silent no-op: a template asking for
/// `separators=(',', ':')` and getting our spaced output would be exactly the invisible byte drift
/// this module exists to prevent.
fn tojson(value: JValue, kwargs: Kwargs) -> Result<String, JError> {
    let indent: Option<usize> = kwargs.get("indent").unwrap_or(None);
    let ensure_ascii: bool = kwargs.get("ensure_ascii").unwrap_or(false);
    kwargs.assert_all_used()?;

    let mut buf = Vec::new();
    let err = |e: String| JError::new(JErrorKind::InvalidOperation, e);
    match indent {
        // `json.dumps(indent=N)` switches the separators to `(',', ': ')` and pretty-prints —
        // which is exactly serde_json's `PrettyFormatter`, including empty `{}` / `[]`.
        Some(n) => {
            let pad = " ".repeat(n);
            let fmt = serde_json::ser::PrettyFormatter::with_indent(pad.as_bytes());
            let mut ser = serde_json::Serializer::with_formatter(&mut buf, fmt);
            value.serialize(&mut ser).map_err(|e| err(e.to_string()))?;
        }
        None => {
            let mut ser = serde_json::Serializer::with_formatter(&mut buf, SpaceFormatter);
            value.serialize(&mut ser).map_err(|e| err(e.to_string()))?;
        }
    }
    let out = String::from_utf8(buf).map_err(|e| err(e.to_string()))?;
    Ok(if ensure_ascii {
        escape_non_ascii(&out)
    } else {
        out
    })
}

/// `json.dumps(ensure_ascii=True)`: every non-ASCII scalar becomes a `\uXXXX` escape, with
/// astral-plane characters written as a UTF-16 surrogate pair. Safe to apply to finished JSON —
/// non-ASCII can only occur inside string literals, every structural byte being ASCII.
fn escape_non_ascii(s: &str) -> String {
    if s.is_ascii() {
        return s.to_string();
    }
    let mut out = String::with_capacity(s.len());
    for c in s.chars() {
        if c.is_ascii() {
            out.push(c);
        } else {
            let mut buf = [0u16; 2];
            for unit in c.encode_utf16(&mut buf) {
                out.push_str(&format!("\\u{unit:04x}"));
            }
        }
    }
    out
}

/// `s | lstrip(chars)` — Python's `str.lstrip`, as a filter. Ported from LiteRT-LM's shim set;
/// jinja2 has no such filter (only `trim`), so this is additive. The method form `s.lstrip(chars)`
/// that templates actually use comes from `pycompat`.
fn lstrip(s: std::borrow::Cow<'_, str>, chars: Option<std::borrow::Cow<'_, str>>) -> String {
    match chars {
        Some(chars) => {
            let set = chars.chars().collect::<Vec<_>>();
            s.trim_start_matches(&set[..]).to_string()
        }
        None => s.trim_start().to_string(),
    }
}

/// `s | rstrip(chars)` — see [`lstrip`].
fn rstrip(s: std::borrow::Cow<'_, str>, chars: Option<std::borrow::Cow<'_, str>>) -> String {
    match chars {
        Some(chars) => {
            let set = chars.chars().collect::<Vec<_>>();
            s.trim_end_matches(&set[..]).to_string()
        }
        None => s.trim_end().to_string(),
    }
}

/// A [`TemplateError`] is always the caller's fault from the server's point of view — a role order
/// the checkpoint rejects, or content in a shape its template cannot express — so it is a 400.
#[cfg(feature = "server")]
impl From<TemplateError> for super::error::ApiError {
    fn from(e: TemplateError) -> Self {
        super::error::ApiError::invalid_request(e.to_string()).with_param("messages")
    }
}

/// Wire messages → the `messages` array a chat template expects.
///
/// Two shape changes the OpenAI wire format forces:
/// - `function.arguments` is a **string** on the wire but every HF template treats it as an
///   object (`tool_call.function.arguments | tojson`), so it is parsed back. Unparseable
///   arguments are passed through as a string rather than dropped — a malformed replay should
///   reach the model as text, not vanish.
/// - Absent optional fields (`tool_call_id`, `name`) are omitted, not nulled: templates test them
///   with `is defined` as often as with `is none`.
///
/// `tools_preamble` is the [`super::tools::tools_system_block`] text, supplied only when the
/// template has no tool grammar of its own ([`TemplateCapabilities::supports_tools`] false). It is
/// folded into the leading system turn, or carried by a synthetic one — the same placement
/// [`super::chat::render_chat_prompt`] uses, so the two paths advertise tools identically.
///
/// `caps` drives one more shape decision: a template with
/// [`TemplateCapabilities::requires_typed_content`] iterates `message.content` as a list of typed
/// blocks, so a plain string yields an EMPTY turn rather than an error. No checkpoint we serve
/// today needs it, but every multimodal template does, and the failure is silent.
#[cfg(feature = "server")]
pub fn wire_messages(
    msgs: &[super::types::ChatMessage],
    tools_preamble: Option<&str>,
    caps: TemplateCapabilities,
) -> serde_json::Value {
    use serde_json::{Map, Value};

    /// `"hi"` or `[{"type": "text", "text": "hi"}]`, per the template's content convention.
    fn content_value(text: &str, typed: bool) -> Value {
        if typed {
            let mut block = Map::new();
            block.insert("type".into(), Value::from("text"));
            block.insert("text".into(), Value::from(text));
            Value::Array(vec![Value::Object(block)])
        } else {
            Value::from(text)
        }
    }
    let typed = caps.requires_typed_content;

    let mut out: Vec<Value> = Vec::with_capacity(msgs.len() + 1);
    let first_is_system = msgs.first().is_some_and(|m| m.role == "system");
    if let Some(block) = tools_preamble
        && !first_is_system
    {
        let mut sys = Map::new();
        sys.insert("role".into(), Value::from("system"));
        sys.insert("content".into(), content_value(block, typed));
        out.push(Value::Object(sys));
    }

    for (i, m) in msgs.iter().enumerate() {
        let mut obj = Map::new();
        obj.insert("role".into(), Value::from(m.role.as_str()));

        let mut content = m.content.clone();
        if i == 0
            && first_is_system
            && let Some(block) = tools_preamble
        {
            if !content.is_empty() {
                content.push_str("\n\n");
            }
            content.push_str(block);
        }
        obj.insert("content".into(), content_value(&content, typed));

        if !m.tool_calls.is_empty() {
            let calls = m
                .tool_calls
                .iter()
                .map(|tc| {
                    let raw = tc.function.arguments.trim();
                    let args = if raw.is_empty() {
                        Value::Object(Map::new())
                    } else {
                        serde_json::from_str(raw).unwrap_or_else(|_| Value::from(raw))
                    };
                    let mut func = Map::new();
                    func.insert("name".into(), Value::from(tc.function.name.as_str()));
                    func.insert("arguments".into(), args);
                    let mut call = Map::new();
                    if let Some(id) = &tc.id {
                        call.insert("id".into(), Value::from(id.as_str()));
                    }
                    call.insert(
                        "type".into(),
                        Value::from(tc.kind.as_deref().unwrap_or("function")),
                    );
                    call.insert("function".into(), Value::Object(func));
                    Value::Object(call)
                })
                .collect::<Vec<_>>();
            obj.insert("tool_calls".into(), Value::Array(calls));
        }
        if let Some(id) = &m.tool_call_id {
            obj.insert("tool_call_id".into(), Value::from(id.as_str()));
        }
        if let Some(name) = &m.name {
            obj.insert("name".into(), Value::from(name.as_str()));
        }
        out.push(Value::Object(obj));
    }
    Value::Array(out)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn render(source: &str, inputs: &RenderInputs) -> Result<String, TemplateError> {
        ChatTemplate::new(source).expect("compile").render(inputs)
    }

    #[test]
    fn renders_a_basic_conversation() {
        let out = render(
            "Hello {{ messages[0].content }}",
            &RenderInputs {
                messages: serde_json::json!([{"role": "user", "content": "World"}]),
                ..Default::default()
            },
        )
        .expect("render");
        assert_eq!(out, "Hello World");
    }

    #[test]
    fn tojson_spacing_matches_python_json_dumps() {
        let mut extra = serde_json::Map::new();
        extra.insert("data".into(), serde_json::json!({"a": 1, "b": [1, 2]}));
        let out = render(
            "{{ data | tojson }}",
            &RenderInputs {
                extra,
                ..Default::default()
            },
        )
        .expect("render");
        assert_eq!(out, r#"{"a": 1, "b": [1, 2]}"#);
    }

    #[test]
    fn tojson_indent_matches_python_pretty_printing() {
        let mut extra = serde_json::Map::new();
        extra.insert("data".into(), serde_json::json!({"a": [1], "e": {}}));
        let out = render(
            "{{ data | tojson(indent=4) }}",
            &RenderInputs {
                extra,
                ..Default::default()
            },
        )
        .expect("render");
        assert_eq!(out, "{\n    \"a\": [\n        1\n    ],\n    \"e\": {}\n}");
    }

    #[test]
    fn tojson_ensure_ascii_escapes_like_cpython() {
        let mut extra = serde_json::Map::new();
        extra.insert("data".into(), serde_json::json!({"s": "é🙂"}));
        // Default is ensure_ascii=False (transformers overrides jinja2 here).
        let raw = render(
            "{{ data | tojson }}",
            &RenderInputs {
                extra: extra.clone(),
                ..Default::default()
            },
        )
        .expect("render");
        assert_eq!(raw, "{\"s\": \"é🙂\"}");
        // Opt in and the astral char becomes a surrogate pair, exactly as json.dumps writes it.
        let escaped = render(
            "{{ data | tojson(ensure_ascii=True) }}",
            &RenderInputs {
                extra,
                ..Default::default()
            },
        )
        .expect("render");
        // Verbatim `json.dumps({"s": "é🙂"}, ensure_ascii=True)` — the astral char written as
        // a UTF-16 surrogate pair, which is what CPython emits.
        assert_eq!(escaped, r#"{"s": "\u00e9\ud83d\ude42"}"#);
    }

    #[test]
    fn an_unsupported_tojson_keyword_is_loud() {
        let mut extra = serde_json::Map::new();
        extra.insert("data".into(), serde_json::json!([1]));
        let err = render(
            "{{ data | tojson(separators=',') }}",
            &RenderInputs {
                extra,
                ..Default::default()
            },
        )
        .expect_err("unknown keyword must fail");
        assert_eq!(err.stage, TemplateStage::Render);
    }

    #[test]
    fn raise_exception_is_an_error_carrying_the_message() {
        let err = render("{{ raise_exception('nope') }}", &RenderInputs::default())
            .expect_err("must fail");
        assert!(err.message.contains("nope"), "{err}");
        assert_eq!(err.stage, TemplateStage::Render);
    }

    #[test]
    fn strftime_now_formats_the_supplied_instant_in_utc() {
        let out = render(
            "{{ strftime_now('%Y-%m-%d') }}",
            &RenderInputs {
                now: Some(1735689600),
                ..Default::default()
            },
        )
        .expect("render");
        assert_eq!(out, "2025-01-01");
    }

    #[test]
    fn lstrip_and_rstrip_take_a_char_set() {
        let out = render(
            "{{ '  bar  '|lstrip }}|{{ '  baz  '|rstrip }}|{{ '1212foo12'|lstrip('12') }}",
            &RenderInputs::default(),
        )
        .expect("render");
        assert_eq!(out, "bar  |  baz|foo12");
    }

    #[test]
    fn python_string_methods_resolve() {
        let out = render(
            "{{ 'a,b'.split(',') | join('|') }}/{{ 'xy'.startswith('x') }}/{{ 'ab'.replace('a','c') }}",
            &RenderInputs::default(),
        )
        .expect("render");
        // `True`, not `true` — jinja2 writes `str(bool)`. See `python_str`.
        assert_eq!(out, "a|b/True/cb");
    }

    /// `{{ dict }}` and `{{ dict | string }}` are `str()` in jinja2, i.e. a Python **repr** with
    /// single quotes — not JSON. SmolLM3 advertises its tools exactly this way, so getting it
    /// wrong silently rewrites every tool prompt for that family.
    #[test]
    fn values_stringify_the_way_cpython_does() {
        let mut extra = serde_json::Map::new();
        extra.insert(
            "d".into(),
            serde_json::json!({"type": "function", "ok": true, "n": null, "xs": [1, "a"]}),
        );
        extra.insert("s".into(), serde_json::json!("plain"));
        extra.insert("q".into(), serde_json::json!(["it's", "say \"hi\""]));
        let out = render(
            "{{ d }}|{{ d | string }}|{{ s }}|{{ true }}|{{ none }}|{{ q }}",
            &RenderInputs {
                extra,
                ..Default::default()
            },
        )
        .expect("render");
        assert_eq!(
            out,
            // `python3 -c "print({'type':'function','ok':True,'n':None,'xs':[1,'a']})"`, then the
            // same through `str()`, then a bare string (no quotes — `str('plain') == 'plain'`),
            // then True/None, then repr's quote-flip when the string already contains `'`.
            "{'type': 'function', 'ok': True, 'n': None, 'xs': [1, 'a']}\
             |{'type': 'function', 'ok': True, 'n': None, 'xs': [1, 'a']}\
             |plain|True|None|[\"it's\", 'say \"hi\"']"
        );
    }

    /// jinja2: an undefined variable is NOT none, and an explicit null IS. The whole tool-aware
    /// family branches on this, so the builtin test must not be remapped to `is_undefined`.
    #[test]
    fn the_none_test_follows_jinja2_semantics() {
        let out = render(
            "{{ tools is none }}/{{ nothing is none }}/{{ nothing is undefined }}",
            &RenderInputs::default(),
        )
        .expect("render");
        assert_eq!(out, "True/False/True");
    }

    /// A template that fails to parse comes back as an error the caller can fall back from.
    #[test]
    fn a_broken_template_fails_to_compile_without_panicking() {
        let err = ChatTemplate::new("{% for x in %}")
            .err()
            .expect("must fail");
        assert_eq!(err.stage, TemplateStage::Compile);
    }
}