# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Declarative configuration for model tool-call and reasoning formats.

Each model format is described by a :class:`ParserEngineConfig` that specifies:

* **terminals** – literal strings or regex patterns that delimit the format
  (e.g. ``<tool_call>``, ``</think>``).
* **token_id_terminals** – terminals that should be matched by token ID
  rather than (or in addition to) text.
* **transitions** – a state machine mapping
  ``(state, terminal) → (new_state, events_to_emit)`` that drives semantic
  event generation during streaming.
* **content_events** – what :class:`EventType` to emit for plain content
  (non-terminal text) in each state.
"""

from __future__ import annotations

from collections.abc import Callable, Mapping
from dataclasses import dataclass, field
from enum import Enum, auto
from functools import cached_property

from vllm.parser.engine.events import EventType

# Maps a raw argument delta to the JSON text it appends to the open trailing
# value, or returns None when the delta may change anything else.
ArgValueExtender = Callable[[str], str | None]


class ParserState(Enum):
    CONTENT = auto()
    REASONING = auto()
    MESSAGE_HEADER = auto()
    TOOL_PREAMBLE = auto()
    TOOL_NAME = auto()
    TOOL_ARGS = auto()
    TOOL_BETWEEN = auto()


@dataclass(frozen=True, slots=True)
class Transition:
    next_state: ParserState
    events: tuple[EventType, ...] = field(default_factory=tuple)
    skip_in_token_id_mode: bool = False


@dataclass(frozen=True)
class ParserEngineConfig:
    """Declarative description of a model's tool-call / reasoning format.

    The engine feeds terminals from the incremental lexer into the
    transition table and emits the corresponding semantic events.
    Content tokens (text between terminals) are classified by the
    current state via ``content_events``.
    """

    name: str

    # A terminal may carry several spellings; the first spelling is the canonical one.
    terminals: Mapping[str, str | tuple[str, ...]] = field(default_factory=dict)

    token_id_terminals: dict[str, str] = field(default_factory=dict)

    transitions: dict[tuple[ParserState, str], Transition] = field(
        default_factory=dict,
    )

    content_events: dict[ParserState, EventType] = field(
        default_factory=lambda: {
            ParserState.CONTENT: EventType.TEXT_CHUNK,
            ParserState.REASONING: EventType.REASONING_CHUNK,
            ParserState.TOOL_NAME: EventType.TOOL_NAME,
            ParserState.TOOL_ARGS: EventType.ARG_VALUE_CHUNK,
        },
    )

    initial_state: ParserState = ParserState.CONTENT

    wait_for_reasoning: bool | None = None

    arg_converter: Callable[[str, bool], str] | None = None

    stream_arg_deltas: bool = True

    tool_args_json: bool = True

    # The tool region body is a single JSON array of ``{"name", "arguments"}``
    # objects (one marker → N calls, no closing terminal). The engine splits
    # the array into one tool call per element; the array's ``]`` ends the
    # region and any trailing text becomes content. Used by Granite.
    tool_call_body_array: bool = False

    arg_structural_chars: frozenset[str] | None = None

    # Called with the raw argument text when its converted JSON ends in an
    # open value; returns an extender for that value, or None. Streaming
    # skips ``arg_converter`` for every delta the extender accepts.
    arg_value_extender: Callable[[str], ArgValueExtender | None] | None = None

    # Special tokens exempt from auto-drop but not state-machine terminals.
    preserve_tokens: frozenset[str] = field(default_factory=frozenset)

    # Special tokens delimiting conversation turns in the prompt only
    # considers reasoning markers after the last boundary token.
    turn_boundary_tokens: frozenset[str] = field(default_factory=frozenset)

    # Prevents trailing-whitespace accumulation across multi-turn conversations.
    strip_trailing_reasoning_whitespace: bool = True

    # Drop content that is entirely whitespace when tool calls follow.
    drop_whitespace_only_content_before_tools: bool = True

    # .strip() content text when tool calls are present.
    strip_content_whitespace_with_tools: bool = True

    # Reject tool calls whose names are absent from the request tools.
    validate_tool_names: bool = False

    def terminal_literal(self, name: str) -> str | None:
        """Canonical spelling of terminal *name*, or ``None`` if undeclared."""
        value = self.terminals.get(name)
        if isinstance(value, tuple):
            return value[0] if value else None
        return value

    @property
    def terminal_literals(self) -> frozenset[str]:
        """Every declared spelling across all terminals."""
        return frozenset(
            lit
            for value in self.terminals.values()
            for lit in ((value,) if isinstance(value, str) else value)
        )

    @cached_property
    def terminal_defs(self):
        from vllm.parser.engine.incremental_lexer import terminals_from_literals

        return terminals_from_literals(self.terminals)

    @cached_property
    def lexer_shape(self):
        from vllm.parser.engine.incremental_lexer import LexerShape

        return LexerShape(self.terminal_defs)