"""Custom vLLM tool-call parser for Mistral 3.x served via the HF tokenizer.

Why this exists
---------------
vLLM's built-in `mistral` parser supports the
`[TOOL_CALLS]<name>[CALL_ID]<id>[ARGS]<json>` output format ONLY when the
tokenizer is a `MistralTokenizer` (mistral-common) of version >= 11. When
the model is served via the standard HF tokenizer (tokenizer_mode=auto) —
which is required for any merged Unsloth/HF Mistral-3.x model that ships
config.json + tokenizer.json instead of params.json — the parser falls
back to a JSON-list regex (`[{...}]`) and silently fails to extract.

This plugin registers a `mistral_hf` parser that handles the actual
output format Mistral-Small-3.x emits on the HF path:

    [TOOL_CALLS]<name>[CALL_ID]<9-char id>?[ARGS]<json>
    [TOOL_CALLS]<name>[ARGS]<json>                       (CALL_ID optional)

…and supports multiple concatenated calls.

Usage on the vLLM CLI:
    vllm serve <model> \
        --enable-auto-tool-choice \
        --tool-parser-plugin /path/to/mistral_hf_tool_parser.py \
        --tool-call-parser mistral_hf

Worker-vllm env equivalents (the worker auto-maps these to AsyncEngineArgs):
    TOOL_PARSER_PLUGIN=/path/to/mistral_hf_tool_parser.py
    TOOL_CALL_PARSER=mistral_hf
    ENABLE_AUTO_TOOL_CHOICE=true
"""

from __future__ import annotations

import json
import random
import re
import string
from collections.abc import Sequence
from typing import Any

try:
    # vLLM <= 0.11.x layout: everything in entrypoints.openai.protocol
    from vllm.entrypoints.openai.protocol import (
        ChatCompletionRequest,
        DeltaFunctionCall,
        DeltaMessage,
        DeltaToolCall,
        ExtractedToolCallInformation,
        FunctionCall,
        ToolCall,
    )
except ImportError:
    # vLLM 0.19.x split protocol.py into per-feature submodules.
    # ChatCompletionRequest lives under chat_completion; the tool-call
    # delta / extraction / streaming types live under engine.
    from vllm.entrypoints.openai.chat_completion.protocol import (
        ChatCompletionRequest,
    )
    from vllm.entrypoints.openai.engine.protocol import (
        DeltaFunctionCall,
        DeltaMessage,
        DeltaToolCall,
        ExtractedToolCallInformation,
        FunctionCall,
        ToolCall,
    )

try:
    # vLLM <= 0.11.x layout
    from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import (
        ToolParser,
        ToolParserManager,
    )
except ImportError:
    # vLLM 0.19.x flattened tool_parsers to a top-level module.
    from vllm.tool_parsers.abstract_tool_parser import (
        ToolParser,
        ToolParserManager,
    )

from vllm.logger import init_logger

logger = init_logger(__name__)

# Matches one tool call:
#   [TOOL_CALLS]<name>[CALL_ID]<id>?[ARGS]<json>
# - name:   letters/digits/_/-
# - id:     optional, alphanumeric (Mistral spec is exactly 9 chars, but we
#           accept any length to be tolerant — the API response just needs
#           SOME id, we generate one if the model emits CALL_ID-less form)
# - args:   non-greedy match up to the next [TOOL_CALLS] or end-of-string,
#           validated via json.loads
_CALL_RE = re.compile(
    r"\[TOOL_CALLS\]"
    r"\s*(?P<name>[A-Za-z0-9_-]+)\s*"
    r"(?:\[CALL_ID\]\s*(?P<id>[A-Za-z0-9]+)\s*)?"
    r"\[ARGS\]\s*(?P<args>\{.*?\})"
    r"(?=\s*\[TOOL_CALLS\]|\s*$)",
    re.DOTALL,
)

_ALPHANUM = string.ascii_letters + string.digits


def _gen_id() -> str:
    return "".join(random.choices(_ALPHANUM, k=9))


def _parse_all(text: str) -> list[dict[str, Any]]:
    """Find every `[TOOL_CALLS]name[CALL_ID]?id[ARGS]{json}` occurrence."""
    calls: list[dict[str, Any]] = []
    for m in _CALL_RE.finditer(text):
        name = m.group("name")
        call_id = m.group("id") or _gen_id()
        raw_args = m.group("args")
        try:
            parsed_args = json.loads(raw_args)
        except json.JSONDecodeError:
            logger.warning(
                "mistral_hf parser: could not JSON-decode args for %s: %r",
                name,
                raw_args[:200],
            )
            continue
        calls.append(
            {
                "id": call_id,
                "name": name,
                "arguments": parsed_args,
            }
        )
    return calls


@ToolParserManager.register_module("mistral_hf")
class MistralHFToolParser(ToolParser):
    """Tool parser for Mistral 3.x served via the HF tokenizer path."""

    BOT_TOKEN = "[TOOL_CALLS]"

    def __init__(self, tokenizer, *args, **kwargs):
        # vLLM 0.11.x ToolParser.__init__ takes (tokenizer); vLLM 0.19.x
        # takes (tokenizer, tools=None). Forward whatever extras the host
        # vLLM passes so this plugin works on either lineage.
        super().__init__(tokenizer, *args, **kwargs)
        # Streaming state
        self.prev_tool_call_arr: list[dict] = []
        self.current_tool_id: int = -1
        self.current_tool_name_sent: bool = False
        self.streamed_args_for_tool: list[str] = []

    def adjust_request(self, request: ChatCompletionRequest) -> ChatCompletionRequest:
        # vLLM strips special tokens from detokenized output by default. We
        # need [TOOL_CALLS] / [ARGS] / [CALL_ID] preserved in `model_output`
        # for our regex to fire. Match the upstream MistralToolParser behavior.
        request = super().adjust_request(request)
        if request.tools and request.tool_choice != "none":
            request.skip_special_tokens = False
        return request

    def extract_tool_calls(
        self,
        model_output: str,
        request: ChatCompletionRequest,
    ) -> ExtractedToolCallInformation:
        if self.BOT_TOKEN not in model_output:
            return ExtractedToolCallInformation(
                tools_called=False,
                tool_calls=[],
                content=model_output,
            )

        try:
            calls = _parse_all(model_output)
            if not calls:
                # Token was present but the format didn't match — surface as
                # plain content so the user still sees the model's text.
                logger.warning(
                    "mistral_hf parser: found [TOOL_CALLS] but no callable "
                    "payload parsed. Returning as content."
                )
                return ExtractedToolCallInformation(
                    tools_called=False,
                    tool_calls=[],
                    content=model_output,
                )

            tool_calls = [
                ToolCall(
                    id=c["id"],
                    type="function",
                    function=FunctionCall(
                        name=c["name"],
                        arguments=json.dumps(c["arguments"], ensure_ascii=False),
                    ),
                )
                for c in calls
            ]
            # Anything before the first [TOOL_CALLS] is preamble content.
            preamble = model_output.split(self.BOT_TOKEN, 1)[0]
            return ExtractedToolCallInformation(
                tools_called=True,
                tool_calls=tool_calls,
                content=preamble if preamble else None,
            )
        except Exception:
            logger.exception(
                "mistral_hf parser: error while extracting tool calls; "
                "returning raw text as content"
            )
            return ExtractedToolCallInformation(
                tools_called=False,
                tool_calls=[],
                content=model_output,
            )

    def extract_tool_calls_streaming(
        self,
        previous_text: str,
        current_text: str,
        delta_text: str,
        previous_token_ids: Sequence[int],
        current_token_ids: Sequence[int],
        delta_token_ids: Sequence[int],
        request: ChatCompletionRequest,
    ) -> DeltaMessage | None:
        # Minimal streaming support: stream non-tool content verbatim; once
        # one or more complete [TOOL_CALLS]<name>[ARGS]<json> payloads are
        # parseable in current_text, emit them in a single DeltaMessage.
        # Each call goes out as one whole DeltaToolCall (id + name + args)
        # rather than streaming arguments piecewise — fine for the small
        # tool-args payloads in this workload.
        if self.BOT_TOKEN not in current_text:
            return DeltaMessage(content=delta_text)

        calls = _parse_all(current_text)
        if len(calls) <= len(self.prev_tool_call_arr):
            # We've seen the BOT token but no new complete call(s) yet —
            # likely still receiving the JSON args. Emit nothing.
            return None

        prev_len = len(self.prev_tool_call_arr)
        new_calls = calls[prev_len:]
        deltas = []
        for offset, c in enumerate(new_calls):
            self.current_tool_id += 1
            args_str = json.dumps(c["arguments"], ensure_ascii=False)
            # DeltaToolCall (NOT ToolCall) — DeltaMessage.tool_calls expects
            # the streaming delta type, and DeltaToolCall requires `index`.
            deltas.append(
                DeltaToolCall(
                    id=c["id"],
                    type="function",
                    index=prev_len + offset,
                    function=DeltaFunctionCall(
                        name=c["name"],
                        arguments=args_str,
                    ),
                )
            )
            # vLLM's finish-reason path reads `streamed_args_for_tool[index]`
            # and `prev_tool_call_arr[index].arguments` to detect any unsent
            # argument tail. Keep both lists in sync with what we emitted,
            # storing arguments as the already-serialized JSON string so
            # the upstream `expected_call.replace(actual_call, ...)` sees
            # equal values and produces an empty remaining_call. Mismatched
            # list lengths cause `IndexError: list index out of range` deep
            # inside serving.py.
            self.streamed_args_for_tool.append(args_str)
        # Store arguments as the JSON string (not the parsed dict) so the
        # upstream code path in serving.py treats it consistently.
        self.prev_tool_call_arr = [
            {
                "name": c["name"],
                "arguments": json.dumps(c["arguments"], ensure_ascii=False),
                "id": c["id"],
            }
            for c in calls
        ]
        return DeltaMessage(tool_calls=deltas)