"""上下文预算守卫：历史过长时折叠中间消息，保护 prompt 预算。"""

from app.ai_provider.types.messages import BaseMessage
from app.application.agent.extensions import (
    AgentContext,
    AgentRunContext,
    ExtensionRegistry,
)
from app.application.agent.python_skills import PythonSkillDefinition

# messages 总字符数软上限；超出则保留首条与末尾 _TAIL_KEEP 条，中间折叠。
_MAX_CONTEXT_CHARS = 12000
_TAIL_KEEP = 8


async def _trim_history(context: AgentContext, _run: AgentRunContext) -> AgentContext:
    messages = context.messages
    total = sum(len(message.content) for message in messages)
    if total <= _MAX_CONTEXT_CHARS or len(messages) <= _TAIL_KEEP + 1:
        return context
    first = messages[0]
    tail = messages[-_TAIL_KEEP:]
    dropped = len(messages) - 1 - _TAIL_KEEP
    # 在首条追加折叠提示（新建对象避免污染原 message）；不在历史中间插入新消息，
    # 以免破坏 provider 对 system 位置或多轮 tool 消息配对的预期。
    folded = BaseMessage(
        content=f"{first.content}\n\n[已省略 {dropped} 条较早历史]",
        role=first.role,
    )
    return AgentContext(
        messages=(folded, *tail),
        tools=context.tools,
        model_hint=context.model_hint,
    )


def _setup(registry: ExtensionRegistry) -> None:
    registry.add_context_transformer(_trim_history)


SKILL = PythonSkillDefinition(
    name="context_budget",
    description="对话历史过长时折叠中间消息，控制 prompt 预算。",
    triggers=("长对话", "context", "历史"),
    instructions="历史较长时优先依赖近期消息，不必回溯很早的内容。",
    setup=_setup,
)