# Copyright 2025 Horizon RL Contributors

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Base classes for tool call parsing."""

from __future__ import annotations

import json
import re
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Callable, ClassVar, TypeVar

if TYPE_CHECKING:
    from transformers import PreTrainedTokenizerBase

# Parser registry - populated by @register_tool_parser decorator
TOOL_PARSER_REGISTRY: dict[str, type[ToolParser]] = {}

T = TypeVar("T", bound="ToolParser")


@dataclass(frozen=True)
class ToolParseResult:
    """A parsed tool call request.

    For successful parses: name and input are populated, raw is None.
    For parse errors: name is extracted or UNKNOWN_NAME, raw contains the unparseable content.
    """

    UNKNOWN_NAME: ClassVar[str] = "unknown_tool"

    id: str
    name: str
    input: dict[str, Any] = field(default_factory=dict)
    raw: str | None = None

    @classmethod
    def from_parse_error(cls, id: str, raw: str, name: str | None = None) -> ToolParseResult:
        """Create an error result for parse failures.

        Args:
            id: Tool call ID.
            raw: The unparseable raw content (fed back to model for self-correction).
            name: Best-effort extracted tool name (defaults to UNKNOWN_NAME).
        """
        return cls(id=id, name=name if name is not None else cls.UNKNOWN_NAME, input={}, raw=raw)

    @property
    def is_error(self) -> bool:
        """Check if this represents a parse error."""
        return self.raw is not None

    @property
    def payload(self) -> str:
        """Get the tool call payload string to pass to the tool executor.

        For successful parses, returns JSON-encoded input.
        For errors, returns the raw content (so model sees its mistake).
        """
        if self.is_error:
            return self.raw or ""
        return json.dumps(self.input)


class ToolParser(ABC):
    """Base class for tool call parsers.

    Subclasses implement `parse` to extract tool calls from model output.
    Only JSONDecodeError is handled; Strands validates arguments downstream.

    Example:
        >>> from strands_sglang import get_tool_parser
        >>> parser = get_tool_parser("hermes")
        >>> results = parser.parse('<tool_call>{"name": "foo", "arguments": {}}</tool_call>')
        >>> print(results[0].name)
        foo
    """

    DEFAULT_TOOL_START_TOKEN = "<tool_call>"
    DEFAULT_TOOL_END_TOKEN = "</tool_call>"
    DEFAULT_THINK_START_TOKEN = "<think>"
    DEFAULT_THINK_END_TOKEN = "</think>"

    def __init__(
        self,
        tool_start_token: str = DEFAULT_TOOL_START_TOKEN,
        tool_end_token: str = DEFAULT_TOOL_END_TOKEN,
        think_start_token: str = DEFAULT_THINK_START_TOKEN,
        think_end_token: str = DEFAULT_THINK_END_TOKEN,
    ) -> None:
        """Initialize the parser with optional custom tokens.

        Args:
            tool_start_token: Opening token for tool calls.
            tool_end_token: Closing token for tool calls.
            think_start_token: Opening token for think blocks.
            think_end_token: Closing token for think blocks.
        """
        self.tool_start_token = tool_start_token
        self.tool_end_token = tool_end_token
        self.think_start_token = think_start_token
        self.think_end_token = think_end_token

        # Pattern to extract tool call content (with whitespace trimming)
        self.tool_pattern = re.compile(
            rf"{re.escape(tool_start_token)}\s*(.*?)\s*{re.escape(tool_end_token)}",
            re.DOTALL,
        )
        # Pattern to remove think blocks (no capture needed)
        self.think_pattern = re.compile(
            rf"{re.escape(think_start_token)}.*?{re.escape(think_end_token)}",
            re.DOTALL,
        )

    def validate_tokenizer(self, tokenizer: PreTrainedTokenizerBase) -> None:
        """Validate that the tokenizer is compatible with this parser.

        Called during `SGLangModel.__init__`. Override to check that the
        tokenizer has required setup (e.g., custom encoding attached).

        Args:
            tokenizer: The tokenizer that will be used for chat template formatting.
        """
        _ = tokenizer  # Used by subclass overrides

    @property
    def message_separator(self) -> str:
        """Separator between messages in the chat template.

        Different tokenizers use different separators between messages.
        This is used during incremental tokenization to ensure the TITO
        trajectory matches what `apply_chat_template` would produce.

        Default is no separator.
        """
        return ""

    @abstractmethod
    def parse(self, text: str) -> list[ToolParseResult]:
        """Parse tool calls from model output text.

        Args:
            text: Model output text.

        Returns:
            List of parsed tool call results.
        """
        ...


def register_tool_parser(name: str) -> Callable[[type[T]], type[T]]:
    """Decorator to register a tool parser class.

    Args:
        name: Registry name for the parser.

    Returns:
        Decorator that registers the class and returns it unchanged.

    Example:
        >>> @register_tool_parser("my_parser")
        ... class MyParser(ToolParser):
        ...     def parse(self, text): ...
    """

    def decorator(cls: type[T]) -> type[T]:
        TOOL_PARSER_REGISTRY[name] = cls
        return cls

    return decorator


def get_tool_parser(name: str, **kwargs: Any) -> ToolParser:
    """Get a tool parser by name.

    Args:
        name: Parser name (e.g., "hermes", "qwen_xml").
        **kwargs: Arguments passed to the parser constructor.

    Returns:
        Instantiated parser.

    Raises:
        KeyError: If parser name is not registered.

    Example:
        >>> parser = get_tool_parser("hermes")
        >>> parser = get_tool_parser("hermes", think_start_token="<reasoning>")
    """
    if name not in TOOL_PARSER_REGISTRY:
        available = ", ".join(sorted(TOOL_PARSER_REGISTRY.keys()))
        raise KeyError(f"Unknown tool parser: {name!r}. Available: {available}")
    return TOOL_PARSER_REGISTRY[name](**kwargs)
