# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from __future__ import annotations

import functools
import json
from dataclasses import replace
from itertools import groupby
from typing import TYPE_CHECKING

from vllm.entrypoints.generate.base.protocol import (
    DeltaFunctionCall,
    DeltaMessage,
    DeltaToolCall,
)
from vllm.parser.engine.events import EventType, SemanticEvent
from vllm.parser.engine.parser_engine import ParserEngine
from vllm.parser.engine.parser_engine_config import (
    ParserEngineConfig,
    ParserState,
    Transition,
)

if TYPE_CHECKING:
    from vllm.entrypoints.openai.chat_completion.protocol import (
        ChatCompletionRequest,
    )
    from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
    from vllm.tokenizers import TokenizerLike
    from vllm.tool_parsers.abstract_tool_parser import Tool

BEGIN_TOOL_REQUESTS = "<|plamo:begin_tool_requests:plamo|>"
END_TOOL_REQUESTS = "<|plamo:end_tool_requests:plamo|>"
BEGIN_TOOL_REQUEST = "<|plamo:begin_tool_request:plamo|>"
END_TOOL_REQUEST = "<|plamo:end_tool_request:plamo|>"
BEGIN_TOOL_NAME = "<|plamo:begin_tool_name:plamo|>"
END_TOOL_NAME = "<|plamo:end_tool_name:plamo|>"
BEGIN_TOOL_ARGUMENTS = (
    "<|plamo:begin_tool_arguments:plamo|><|plamo:constrain|>json<|plamo:msg|>"
)
END_TOOL_ARGUMENTS = "<|plamo:end_tool_arguments:plamo|>"
BEGIN_THINK = "<|plamo:begin_think:plamo|>"
END_THINK = "<|plamo:end_think:plamo|>"
EOT = "<|plamo:tag|>"

# PLaMo3 control markers are split into multiple tokenizer tokens.  These
# atomic pieces must not be treated as unrelated special tokens by the
# parser-engine drop-token handling.
PLAMO_MARKER_TOKENS = frozenset(
    {
        "<|plamo:begin_",
        "<|plamo:end_",
        ":plamo|>",
    }
)


@functools.cache
def plamo3_config(thinking: bool = True) -> ParserEngineConfig:
    if thinking:
        reasoning_terminals = {
            "THINK_START": BEGIN_THINK,
            "THINK_END": END_THINK,
        }
        reasoning_transitions = {
            (ParserState.REASONING, "THINK_START"): Transition(
                ParserState.REASONING,
            ),
            (ParserState.CONTENT, "THINK_START"): Transition(
                ParserState.REASONING,
                (EventType.REASONING_START,),
            ),
            (ParserState.REASONING, "THINK_END"): Transition(
                ParserState.CONTENT,
                (EventType.REASONING_END,),
            ),
            (ParserState.CONTENT, "THINK_END"): Transition(
                ParserState.CONTENT,
                (),
            ),
            (ParserState.REASONING, "TOOL_REQUESTS_START"): Transition(
                ParserState.TOOL_PREAMBLE,
                (EventType.REASONING_END,),
            ),
        }
        initial_state = ParserState.REASONING
    else:
        reasoning_terminals = {}
        reasoning_transitions = {}
        initial_state = ParserState.CONTENT

    return ParserEngineConfig(
        name="plamo3",
        initial_state=initial_state,
        terminals={
            **reasoning_terminals,
            "TOOL_REQUESTS_START": BEGIN_TOOL_REQUESTS,
            "TOOL_REQUESTS_END": END_TOOL_REQUESTS,
            "TOOL_REQUEST_START": BEGIN_TOOL_REQUEST,
            "TOOL_REQUEST_END": END_TOOL_REQUEST,
            "TOOL_NAME_START": BEGIN_TOOL_NAME,
            "TOOL_NAME_END": END_TOOL_NAME,
            "TOOL_ARGS_START": BEGIN_TOOL_ARGUMENTS,
            "TOOL_ARGS_END": END_TOOL_ARGUMENTS,
            "EOT": EOT,
        },
        transitions={
            **reasoning_transitions,
            (ParserState.CONTENT, "TOOL_REQUESTS_START"): Transition(
                ParserState.TOOL_PREAMBLE,
            ),
            (ParserState.TOOL_PREAMBLE, "TOOL_REQUEST_START"): Transition(
                ParserState.TOOL_PREAMBLE,
            ),
            (ParserState.TOOL_PREAMBLE, "TOOL_NAME_START"): Transition(
                ParserState.TOOL_NAME,
                (EventType.TOOL_CALL_START,),
            ),
            (ParserState.TOOL_NAME, "TOOL_NAME_END"): Transition(
                ParserState.TOOL_PREAMBLE,
            ),
            (ParserState.TOOL_PREAMBLE, "TOOL_ARGS_START"): Transition(
                ParserState.TOOL_ARGS,
            ),
            (ParserState.TOOL_ARGS, "TOOL_ARGS_END"): Transition(
                ParserState.TOOL_BETWEEN,
            ),
            (ParserState.TOOL_BETWEEN, "TOOL_REQUEST_END"): Transition(
                ParserState.TOOL_PREAMBLE,
                (EventType.TOOL_CALL_END,),
            ),
            (ParserState.TOOL_PREAMBLE, "TOOL_REQUESTS_END"): Transition(
                ParserState.CONTENT,
            ),
            # EOT is normally consumed by generation stop handling.  Keep it
            # in the grammar as well so text-only parsing and explicitly
            # supplied stop tokens cannot leak it into user-visible output.
            (ParserState.REASONING, "EOT"): Transition(
                ParserState.CONTENT,
                (EventType.REASONING_END,),
            ),
            (ParserState.CONTENT, "EOT"): Transition(ParserState.CONTENT),
            (ParserState.TOOL_PREAMBLE, "EOT"): Transition(ParserState.CONTENT),
            (ParserState.TOOL_NAME, "EOT"): Transition(ParserState.CONTENT),
            (ParserState.TOOL_ARGS, "EOT"): Transition(ParserState.CONTENT),
            (ParserState.TOOL_BETWEEN, "EOT"): Transition(ParserState.CONTENT),
        },
        token_id_terminals={"EOT": EOT, **reasoning_terminals},
        drop_whitespace_only_content_before_tools=True,
        stream_arg_deltas=True,
        strip_content_whitespace_with_tools=False,
        strip_trailing_reasoning_whitespace=False,
        preserve_tokens=PLAMO_MARKER_TOKENS,
        tool_args_json=True,
    )


class Plamo3Parser(ParserEngine):
    """PLaMo3 reasoning and tool-call parser backed by ParserEngine."""

    def __init__(
        self,
        tokenizer: TokenizerLike,
        tools: list[Tool] | None = None,
        **kwargs,
    ) -> None:
        # Configure PLaMo parsing and cache its multi-token reasoning markers.
        chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
        self.thinking_enabled = chat_kwargs.get("enable_thinking", True)
        kwargs.setdefault(
            "parser_engine_config",
            plamo3_config(thinking=self.thinking_enabled),
        )
        super().__init__(tokenizer, tools, **kwargs)

        self._reasoning_start_token_id_sequence: list[int] = list(
            tokenizer.encode(BEGIN_THINK, add_special_tokens=False)
        )
        self._reasoning_end_token_id_sequence: list[int] = list(
            tokenizer.encode(END_THINK, add_special_tokens=False)
        )
        # Token-ID sequences of every terminal literal, for token-level
        # unfinished-marker detection.
        self._terminal_token_id_sequences: dict[str, tuple[int, ...]] = {
            literal: tuple(tokenizer.encode(literal, add_special_tokens=False))
            for literal in self.parser_engine_config.terminal_literals
        }

    def is_reasoning_end(self, input_ids: list[int]) -> bool:
        # Detect PLaMo reasoning boundaries that span multiple token IDs.
        if not self.thinking_enabled:
            return True

        start_ids = self._reasoning_start_token_id_sequence
        end_ids = self._reasoning_end_token_id_sequence
        for i in range(len(input_ids) - 1, -1, -1):
            if input_ids[i : i + len(end_ids)] == end_ids:
                return True
            if input_ids[i : i + len(start_ids)] == start_ids:
                return False
            if input_ids[i] == self.vocab.get(EOT):
                return False
        return False

    def extract_content_ids(self, input_ids: list[int]) -> list[int]:
        # Extract content after PLaMo's multi-token reasoning end marker.
        if not self.thinking_enabled:
            return input_ids
        end_ids = self._reasoning_end_token_id_sequence
        for i in range(len(input_ids) - len(end_ids), -1, -1):
            if input_ids[i : i + len(end_ids)] == end_ids:
                return input_ids[i + len(end_ids) :]
        return input_ids

    def _strip_unfinished_marker(self, value: str | None) -> str | None:
        if not value:
            return value
        marker_start = value.rfind("<|plamo:")
        if marker_start < 0:
            return value or None
        suffix = value[marker_start:]
        # Strip a trailing partial marker at the token-ID level.
        suffix_ids = tuple(
            self.model_tokenizer.encode(suffix, add_special_tokens=False)
        )
        if any(
            len(seq) > len(suffix_ids) and seq[: len(suffix_ids)] == suffix_ids
            for seq in self._terminal_token_id_sequences.values()
        ):
            return value[:marker_start] or None
        return value or None

    def _coalesce_finished_tool_events(
        self, events: list[SemanticEvent]
    ) -> list[SemanticEvent]:
        combined = []
        for (kind, _), group in groupby(
            events, key=lambda event: (event.type, event.tool_index)
        ):
            chunks = list(group)
            if kind not in (EventType.TOOL_NAME, EventType.ARG_VALUE_CHUNK):
                combined.extend(chunks)
                continue
            if kind == EventType.TOOL_NAME and chunks[0].tool_index < 0:
                continue
            value = "".join(chunk.value for chunk in chunks)
            if kind == EventType.TOOL_NAME:
                value = self._strip_unfinished_marker(value) or ""
            if value:
                combined.append(replace(chunks[0], value=value))
        return combined

    def _events_to_delta(
        self,
        events: list[SemanticEvent],
        finished: bool = False,
    ) -> DeltaMessage | None:
        # Coalesce flushed tool fragments before stripping truncated markers.
        if not finished:
            return super()._events_to_delta(events)

        events = self._coalesce_finished_tool_events(events)
        delta = super()._events_to_delta(events, finished=True)
        if delta is None:
            return None

        delta.reasoning = self._strip_unfinished_marker(delta.reasoning)
        if not self.skip_tool_parsing:
            delta.content = self._strip_unfinished_marker(delta.content)
        # The flush emits the remaining (never-streamed) tool-call
        # arguments; an empty remainder must match non-streaming
        # extraction (_extract_args_json), which normalizes it to '{}'.
        for tc in delta.tool_calls or []:
            if tc.function is not None and not (tc.function.arguments or "").strip():
                tc.function.arguments = "{}"
        if delta.reasoning is None and delta.content is None and not delta.tool_calls:
            return None
        return delta

    def _handle_arg_chunk(
        self,
        event: SemanticEvent,
        deltas: list[DeltaToolCall],
    ) -> None:
        # Strip truncated marker suffixes and preserve the first argument delta.
        idx = event.tool_index
        slot = self._tool_slots[idx]
        if (marker_pos := event.value.rfind("<|plamo:")) >= 0:
            stripped = event.value[:marker_pos]
            try:
                json.loads(slot.args + stripped)
            except ValueError:
                pass
            else:
                event = replace(event, value=stripped)
        name_sent_before = slot.name_sent
        super()._handle_arg_chunk(event, deltas)
        if event.value and not name_sent_before and slot.name_sent:
            deltas.append(
                DeltaToolCall(
                    index=idx,
                    function=DeltaFunctionCall(arguments=event.value),
                )
            )

    def _extract_args_json(self, raw_args: str, func_name: str) -> str:
        # Return the raw arguments as JSON
        # since PLaMo3 generates function names and arguments separately.
        return raw_args.strip() or "{}"

    def extract_reasoning(
        self,
        model_output: str,
        request: ChatCompletionRequest | ResponsesRequest,
    ) -> tuple[str | None, str | None]:
        reasoning, content = super().extract_reasoning(model_output, request)
        return self._strip_unfinished_marker(reasoning), content