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

from __future__ import annotations

import json
from collections.abc import Sequence
from dataclasses import dataclass, replace
from enum import Enum, auto
from typing import TYPE_CHECKING, NamedTuple

from openai_harmony import HarmonyError, Message, Role
from xgrammar import StructuralTag
from xgrammar.openai_tool_call_schema import BuiltinToolParam, FunctionToolParam
from xgrammar.structural_tag import (
    AnyTextFormat,
    ConstStringFormat,
    Format,
    GrammarFormat,
    JSONSchemaFormat,
    OptionalFormat,
    OrFormat,
    RegexFormat,
    SequenceFormat,
    TagFormat,
    TriggeredTagsFormat,
)

from vllm.entrypoints.chat_utils import make_tool_call_id
from vllm.entrypoints.generate.base.protocol import (
    DeltaFunctionCall,
    DeltaMessage,
    DeltaToolCall,
    FunctionCall,
)
from vllm.entrypoints.openai.chat_completion.protocol import (
    ChatCompletionRequest,
)
from vllm.entrypoints.openai.parser.harmony_utils import (
    extract_function_from_recipient,
    get_streamable_parser_for_assistant,
    is_function_recipient,
)
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.logger import init_logger
from vllm.parser.abstract_parser import DelegatingParser
from vllm.reasoning.gptoss_reasoning_parser import GptOssReasoningParser
from vllm.sampling_params import StructuredOutputsParams
from vllm.tool_parsers.gptoss_tool_parser import GptOssToolParser
from vllm.tool_parsers.structural_tag_registry import (
    SimplifiedToolChoice,
    get_function_parameters,
    register_vllm_structural_tag,
)

if TYPE_CHECKING:
    from openai_harmony import Message, StreamableParser


logger = init_logger(__name__)


class _SegmentType(Enum):
    TOOL = auto()
    REASONING = auto()
    CONTENT = auto()
    IGNORE = auto()

    @staticmethod
    def from_channel_and_recipient(
        channel: str | None, recipient: str | None
    ) -> _SegmentType:
        if recipient and is_function_recipient(recipient):
            return _SegmentType.TOOL
        if channel == "analysis":
            return _SegmentType.REASONING
        if channel == "final" or (channel == "commentary" and recipient is None):
            return _SegmentType.CONTENT
        return _SegmentType.IGNORE


class Segment(NamedTuple):
    channel: str | None
    recipient: str | None
    delta: str
    completed_message: Message | None = None


@dataclass
class ChunkResult:
    segments: list[Segment]
    reasoning_token_count: int


class HarmonyParser(DelegatingParser):
    def __init__(self, tokenizer, tools=None, *args, **kwargs):
        super().__init__(tokenizer, tools, *args, **kwargs)

        if self.reasoning_parser and not isinstance(
            self.reasoning_parser, GptOssReasoningParser
        ):
            raise ValueError(
                "Harmony requires GptOssReasoningParser, "
                f"got {self.reasoning_parser.__class__.__name__}."
            )

        if self.tool_parser and not isinstance(self.tool_parser, GptOssToolParser):
            raise ValueError(
                "Harmony requires GptOssToolParser, "
                f"got {self.tool_parser.__class__.__name__}."
            )

        self._parser: StreamableParser | None = None
        self._next_tool_call_index = 0
        self._num_processed_messages = 0

        # For error recovery
        self._current_message_tokens: list[int] = []

    @property
    def _harmony_parser(self) -> StreamableParser:
        """Lazily initializes the Harmony parser."""
        if self._parser is None:
            self._parser = get_streamable_parser_for_assistant()
        return self._parser

    def _poll_completed_message(self) -> Message | None:
        messages = self._harmony_parser.messages
        if len(messages) <= self._num_processed_messages:
            return None
        msg = messages[self._num_processed_messages]
        msg.recipient = self._normalize_recipient(msg.recipient)
        self._num_processed_messages += 1
        return msg

    def flush(self) -> list[Segment]:
        segments: list[Segment] = []
        try:
            self._harmony_parser.process_eos()
            msg = self._poll_completed_message()
        except HarmonyError:
            logger.warning(
                "Harmony parser ended in a non-terminal state; returning the "
                "recovered raw output."
            )

            final_channel = "final"
            text = self.model_tokenizer.decode(self._current_message_tokens)
            segments.append(
                Segment(
                    channel=final_channel,
                    recipient=None,
                    delta=text,
                    completed_message=None,
                )
            )
            msg = Message.from_role_and_content(Role.ASSISTANT, text).with_channel(
                final_channel
            )

        # Reset to the initial assistant-parser state for the next turn.
        self._parser = None
        self._num_processed_messages = 0
        self._current_message_tokens.clear()

        if msg is None:
            return segments

        segments.append(
            Segment(
                channel=msg.channel,
                recipient=msg.recipient,
                delta="",
                completed_message=msg,
            )
        )
        return segments

    def parse(
        self,
        model_output: str,
        request: ChatCompletionRequest | ResponsesRequest,
        enable_auto_tools: bool = False,
        model_output_token_ids: Sequence[int] = (),
    ) -> tuple[str | None, str | None, list[FunctionCall] | None]:
        """Parse Harmony output from token IDs.

        Tool calls are always extracted regardless of ``enable_auto_tools``.
        Callers must decide whether to surface them.
        """
        result = self.process_chunk(model_output_token_ids)
        flushed_segments = self.flush()
        if flushed_segments:
            result.segments.extend(flushed_segments)

        reasoning_parts: list[str] = []
        content_parts: list[str] = []
        tool_calls: list[FunctionCall] = []

        for segment in result.segments:
            msg = segment.completed_message
            if msg is None:
                continue
            if msg.author.role != "assistant" or not msg.content:
                continue
            text = msg.content[0].text
            segment_type = _SegmentType.from_channel_and_recipient(
                msg.channel, msg.recipient
            )
            match segment_type:
                case _SegmentType.REASONING if self.reasoning_parser and text:
                    reasoning_parts.append(text)
                case _SegmentType.CONTENT if text:
                    content_parts.append(text)
                case _SegmentType.TOOL if self.tool_parser:
                    recipient = msg.recipient
                    content_type = msg.content_type
                    assert recipient is not None
                    if content_type is not None and "json" not in content_type:
                        arguments = text
                    else:
                        try:
                            arguments = json.dumps(json.loads(text))
                        except json.JSONDecodeError:
                            arguments = text
                    tool_calls.append(
                        FunctionCall(
                            name=extract_function_from_recipient(recipient),
                            arguments=arguments,
                        )
                    )

        reasoning = "\n".join(reasoning_parts) or None
        content = "\n".join(content_parts) or None
        return reasoning, content, tool_calls or None

    def parse_delta(
        self,
        delta_text: str,
        delta_token_ids: list[int],
        request: ChatCompletionRequest | ResponsesRequest,
        prompt_token_ids: list[int] | None = None,
        *,
        finished: bool,
    ) -> DeltaMessage | None:
        prev_recipient = self._normalize_recipient(
            self._harmony_parser.current_recipient
        )
        result = self.process_chunk(delta_token_ids)
        if finished:
            flushed_segments = self.flush()
            if flushed_segments:
                result.segments.extend(flushed_segments)
        combined_content = ""
        combined_reasoning = ""
        tool_messages: list[DeltaToolCall] = []

        for segment in result.segments:
            if segment.completed_message is not None:
                prev_recipient = None
                continue

            segment_type = _SegmentType.from_channel_and_recipient(
                segment.channel, segment.recipient
            )
            match segment_type:
                case _SegmentType.REASONING if self.reasoning_parser:
                    combined_reasoning += segment.delta
                case _SegmentType.CONTENT:
                    combined_content += segment.delta
                case _SegmentType.TOOL if self.tool_parser:
                    assert segment.recipient is not None
                    if prev_recipient != segment.recipient:
                        tool_name = extract_function_from_recipient(segment.recipient)
                        tool_messages.append(
                            DeltaToolCall(
                                # HarmonyParser does not use _stream_state;
                                # "random" tool_call_id_type is always used
                                id=make_tool_call_id(),
                                type="function",
                                function=DeltaFunctionCall(
                                    name=tool_name,
                                    arguments=segment.delta,
                                ),
                                index=self._next_tool_call_index,
                            )
                        )
                        self._next_tool_call_index += 1
                        prev_recipient = segment.recipient
                    elif segment.delta:
                        idx = self._next_tool_call_index - 1
                        if tool_messages:
                            tool_msg = tool_messages[-1]
                            assert tool_msg.index == idx
                            fn = tool_msg.function
                            assert fn is not None and fn.arguments is not None
                            fn.arguments += segment.delta
                        else:
                            tool_messages.append(
                                DeltaToolCall(
                                    index=idx,
                                    function=DeltaFunctionCall(arguments=segment.delta),
                                )
                            )

        if finished:
            self._next_tool_call_index = 0

        if not combined_content and not combined_reasoning and not tool_messages:
            return None

        delta_message = DeltaMessage()
        if combined_content:
            delta_message.content = combined_content
        if combined_reasoning:
            delta_message.reasoning = combined_reasoning
        if tool_messages:
            delta_message.tool_calls = tool_messages

        # Suppress reasoning deltas if not requested
        if delta_message and not request.include_reasoning:
            delta_message.reasoning = None

            # If only reasoning was in the message (no content, no tool_calls)
            # skip emitting entirely
            if not delta_message.content and not delta_message.tool_calls:
                return None

        return delta_message

    def process_chunk(self, token_ids: Sequence[int]) -> ChunkResult:
        if not token_ids:
            return ChunkResult(segments=[], reasoning_token_count=0)

        segments: list[Segment] = []
        reasoning_token_count = 0
        for token_id in token_ids:
            self._harmony_parser.process(token_id)
            channel = self._harmony_parser.current_channel
            recipient = self._normalize_recipient(
                self._harmony_parser.current_recipient
            )
            delta = self._harmony_parser.last_content_delta or ""
            completed_message = self._poll_completed_message()

            if completed_message is not None:
                self._current_message_tokens.clear()
            else:
                self._current_message_tokens.append(token_id)

            if channel == "analysis" or (
                channel == "commentary" and recipient is not None
            ):
                reasoning_token_count += 1

            segments.append(
                Segment(
                    channel=channel,
                    recipient=recipient,
                    delta=delta,
                    completed_message=completed_message,
                )
            )

            # TODO: Optionally merge and suppress empty Segments

        return ChunkResult(
            segments=segments,
            reasoning_token_count=reasoning_token_count,
        )

    def adjust_request(
        self, request: ChatCompletionRequest | ResponsesRequest
    ) -> ChatCompletionRequest | ResponsesRequest:
        request = _adjust_output_format(request)
        return super().adjust_request(request)

    @staticmethod
    def _normalize_recipient(recipient: str | None) -> str | None:
        """Remove constrained formats misparsed into recipients by older Harmony."""
        if recipient is None:
            return None

        constrain_index = recipient.find("<|constrain|>")
        if constrain_index == -1:
            return recipient
        return recipient[:constrain_index].rstrip() or None


_JSON_CONSTRAINS = [" json", " <|constrain|>json"]
# Harmomy's stop tokens are <|return|>, <|call|>, <|endoftext|>
# They are represented as "" since xgrammar disallows stop tokens while
# under constraints, leading to bad or infinite generations.
_END_TAG = ["<|end|>", ""]
_FINAL_BEGIN = "<|channel|>final{constrain}<|message|>"
_TOOL_CALL_CHANNELS = [
    "<|channel|>commentary",
    "<|channel|>analysis",
    "<|channel|>final",
]
_FUNCTION_CALL_BEGINS = [
    " to=functions.{name}{channel}{constrain}<|message|>",
    "{channel} to=functions.{name}{constrain}<|message|>",
]
_JSON_CONTENT = JSONSchemaFormat(json_schema={"type": "object"})
_ANY_CONTENT = AnyTextFormat()


def _assemble_tag(
    allow_analysis: bool, allow_commentary: bool, content: Format
) -> StructuralTag:
    tags = []
    if allow_analysis:
        analysis_tag = OptionalFormat(
            content=SequenceFormat(
                elements=[
                    TagFormat(
                        begin="<|channel|>analysis<|message|>",
                        content=_ANY_CONTENT,
                        end="<|end|>",
                    ),
                    ConstStringFormat(value="<|start|>assistant"),
                ]
            )
        )
        tags.append(analysis_tag)

    if allow_commentary:
        commentary_tag = OptionalFormat(
            content=SequenceFormat(
                elements=[
                    TagFormat(
                        begin="<|channel|>commentary<|message|>",
                        content=_ANY_CONTENT,
                        end="<|end|>",
                    ),
                    ConstStringFormat(value="<|start|>assistant"),
                ]
            )
        )
        tags.append(commentary_tag)

    tags.append(content)

    return StructuralTag(format=SequenceFormat(elements=tags))


@register_vllm_structural_tag("harmony")
def get_harmony_structural_tag(
    tools: list[FunctionToolParam],
    builtin_tools: list[BuiltinToolParam],
    tool_choice: SimplifiedToolChoice,
    reasoning: bool,
    token_suffix: str = "",
) -> StructuralTag:
    # reasoning always enabled for Harmony
    del reasoning, token_suffix

    if builtin_tools:
        # Fallback for built-in tools
        tags = [
            TagFormat(
                begin="to=",
                content=AnyTextFormat(excludes=["<|start|>"]),
                end=_END_TAG,
            )
        ]
        tags.extend(
            TagFormat(
                begin=channel + " to=",
                content=AnyTextFormat(excludes=["<|start|>", "<|channel|>"]),
                end=_END_TAG,
            )
            for channel in _TOOL_CALL_CHANNELS
        )
    else:
        tags = [
            TagFormat(
                begin=pattern.format(
                    name=tool.function.name, channel=channel, constrain=constrain
                ),
                content=JSONSchemaFormat(
                    json_schema=get_function_parameters(tool.function)
                ),
                end=_END_TAG,
            )
            for tool in tools
            for pattern in _FUNCTION_CALL_BEGINS
            for constrain in _JSON_CONSTRAINS
            for channel in _TOOL_CALL_CHANNELS
        ]

    if tool_choice == "auto":
        tags.extend(
            TagFormat(
                begin=_FINAL_BEGIN.format(constrain=constrain),
                content=_ANY_CONTENT,
                end=_END_TAG,
            )
            for constrain in _JSON_CONSTRAINS + [""]
        )

    return _assemble_tag(
        allow_analysis=True, allow_commentary=True, content=OrFormat(elements=tags)
    )


def _params_to_final_content(params: StructuredOutputsParams) -> Format | None:
    """Map StructuredOutputsParams in a XGrammar Format."""
    if params.json_object:
        return _JSON_CONTENT
    if params.json is not None:
        schema = params.json
        if isinstance(schema, str):
            schema = json.loads(schema)
        return JSONSchemaFormat(json_schema=schema)
    if params.regex is not None:
        return RegexFormat(pattern=params.regex)
    if params.choice is not None:
        return OrFormat(
            elements=[ConstStringFormat(value=choice) for choice in params.choice]
        )
    if params.grammar is not None:
        return GrammarFormat(grammar=params.grammar)
    if params.structural_tag is not None:
        s_tag = json.loads(params.structural_tag)
        if "structures" in s_tag:
            # LegacyStructuralTagResponseFormat
            return TriggeredTagsFormat(
                triggers=s_tag["triggers"],
                tags=[
                    TagFormat(
                        begin=structure["begin"],
                        content=JSONSchemaFormat(json_schema=structure["schema"]),
                        end=structure["end"],
                    )
                    for structure in s_tag["structures"]
                ],
            )
        # StructuralTagResponseFormat
        return StructuralTag.model_validate(s_tag).format
    return None


def _adjust_output_format(
    request: ChatCompletionRequest | ResponsesRequest,
) -> ChatCompletionRequest | ResponsesRequest:
    """Canonicalize request constraints into a reasoning-aware StructuralTag."""
    params = request.extract_structured_outputs()
    if params is None:
        return request

    final_content = _params_to_final_content(params)
    if final_content is None:
        return request

    if isinstance(final_content, JSONSchemaFormat):
        content = OrFormat(
            elements=[
                TagFormat(
                    begin=_FINAL_BEGIN.format(constrain=constrain),
                    content=final_content,
                    end=_END_TAG,
                )
                for constrain in _JSON_CONSTRAINS
            ]
        )
    else:
        content = TagFormat(
            begin=_FINAL_BEGIN.format(constrain=""), content=final_content, end=_END_TAG
        )
    structural_tag = _assemble_tag(
        allow_analysis=True, allow_commentary=False, content=content
    )

    request.structured_outputs = replace(
        params,
        json=None,
        regex=None,
        choice=None,
        grammar=None,
        json_object=None,
        structural_tag=json.dumps(structural_tag.model_dump()),
    )
    if isinstance(request, ResponsesRequest):
        request.text = None
    else:
        request.response_format = None
    return request