"""Reject runaway argument streams for small, built-in device actions.

This is a transport safety bound, not a model output-token budget. Free-text
tools (input_text, save_note, reports and plugin calls) are deliberately excluded.
Only lengths and tool names are retained; argument bodies may contain user data.
"""

from dataclasses import dataclass


SMALL_ACTION_ARGUMENT_LIMIT = 16 * 1024
_SMALL_ACTIONS = frozenset(
    {"click", "long_press", "swipe", "press_key", "manage_app", "wait_for_delay"}
)


class ToolArgumentOverflow(ValueError):
    def __init__(self, name: str, characters: int):
        self.details = {
            "reason": "tool_arguments_too_large",
            "tool": name,
            "argument_chars": characters,
            "limit_chars": SMALL_ACTION_ARGUMENT_LIMIT,
        }
        super().__init__(
            f"收到的 {name} 参数流超出安全限制（{characters} 字符，"
            f"保护上限 {SMALL_ACTION_ARGUMENT_LIMIT}）；已关闭输出流，未执行本次动作。"
        )


@dataclass
class _ToolStream:
    name: str | None = None
    characters: int = 0


class ToolArgumentGuard:
    def __init__(self):
        self.calls: dict[int | str, _ToolStream] = {}

    def observe(self, chunk) -> None:
        for tool in getattr(chunk, "tool_call_chunks", None) or []:
            index = tool.get("index")
            key = index if index is not None else tool.get("id") or "unknown"
            state = self.calls.setdefault(key, _ToolStream())
            if isinstance(tool.get("name"), str) and tool["name"]:
                state.name = tool["name"]
            arguments = tool.get("args")
            if isinstance(arguments, str):
                state.characters += len(arguments)
            if state.name in _SMALL_ACTIONS and state.characters > SMALL_ACTION_ARGUMENT_LIMIT:
                raise ToolArgumentOverflow(state.name, state.characters)