"""Small, pinned inference projections for the three released model families.

The intent projection mirrors tool-intent-align's evidence-candidate-v3-full-schema
serializer. It takes the definition of the called tool alone rather than the
catalogue it came from, because that is all the projection ever read.
The output prompt follows the released encoder checkpoint's
instruction/separator/data format.
Keep the literal formatting stable and update the comparison tests before changing it.
"""

from __future__ import annotations

import json

ROLE_LABELS = {
    "system": "System",
    "user": "User",
    "assistant": "Assistant",
    "tool": "Tool",
}

def canonical_json(value):
    return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))


def render_intent(context, tool, call):
    """Render the called tool, preserving a missing schema as JSON null."""
    if tool["name"] != call["name"]:
        raise ValueError(
            f"The tool definition is for {tool['name']!r} but the call is to {call['name']!r}."
        )
    lines = []
    for turn in context:
        role = ROLE_LABELS[turn["role"]]
        parts = []
        if turn["content"] is not None:
            parts.append(f"{role}: {turn['content']}")
        for earlier_call in turn["tool_calls"] or ():
            parts.append(f"{role} tool call: {canonical_json(earlier_call)}")
        lines.append("\n".join(parts) or f"{role}:")
    rendered_context = "\n".join(lines)
    rendered_tool = "\n".join(
        (
            "Called tool:",
            f"name: {tool['name']}",
            f"description: {tool.get('description') or ''}",
            f"parameters: {canonical_json(tool.get('parameters'))}",
        )
    )
    return f"Context:\n{rendered_context}\n\n{rendered_tool}", canonical_json(call)


def render_output_filter(instruction, tool_output):
    prefix = f"{instruction} <|end_of_instruction|> "
    return prefix + tool_output, len(prefix), len(prefix) + len(tool_output)


def render_anomaly(task, tool_calls):
    calls = []
    for call in tool_calls:
        arguments = json.dumps(call.get("args", {}), ensure_ascii=False, default=str)[:200]
        calls.append(f"{call['name']}({arguments})")
    return f"TASK: {task or ''} CALLS: " + " ".join(calls)


def merge_spans(spans):
    merged = []
    for start, end in sorted(spans):
        if start >= end:
            continue
        if merged and start <= merged[-1][1]:
            merged[-1][1] = max(merged[-1][1], end)
        else:
            merged.append([start, end])
    return merged


def output_spans(probabilities, offsets, data_start, data_end, threshold):
    """Map DROP token probabilities onto character spans in the original output."""
    spans = []
    for probability, (token_start, token_end) in zip(probabilities, offsets, strict=True):
        if probability < threshold or token_start == token_end:
            continue
        start = max(token_start, data_start)
        end = min(token_end, data_end)
        if start < end:
            spans.append((start - data_start, end - data_start))
    return merge_spans(spans)


def remove_spans(text, spans):
    parts = []
    cursor = 0
    for start, end in spans:
        parts.append(text[cursor:start])
        cursor = end
    parts.append(text[cursor:])
    return "".join(parts)