# -*- coding: utf-8 -*-
"""e-260930031800: хард-кап тул-коллов в воркере (hermetic, без LLM).

Запуск: ./venv/bin/python tests/test_tool_budget.py

Проверяет: _wire_tool_budget отменяет вызовы после N реальных
исполнений (cancel_tool в BeforeToolCallEvent), останавливает цикл
(end_turn в AfterToolsEvent), счётчик растёт, лимит читается из env.
"""
import os
import sys

sys.path.insert(0, "/root/simbiot-harness")

from harness import _llm_worker

fails = []


def check(name, cond, extra=""):
    print(("PASS" if cond else "FAIL"), name, ("| " + extra) if not cond else "")
    if not cond:
        fails.append(name)


class _FakeBeforeEvent:
    def __init__(self):
        self.cancel_tool = False


class _FakeAfterEvent:
    def __init__(self):
        self.end_turn = False


class _FakeAgent:
    def __init__(self):
        self.hooks = []

    def add_hook(self, fn, ev_type):
        self.hooks.append((fn, ev_type))


def _budget_hooks(agent):
    """(before_hook, after_hook) из навешанных на фейковый агент."""
    from strands.hooks import BeforeToolCallEvent, AfterToolsEvent
    before = after = None
    for fn, ev in agent.hooks:
        if ev is BeforeToolCallEvent:
            before = fn
        elif ev is AfterToolsEvent:
            after = fn
    return before, after


# ── 1. хук отменяет вызовы после лимита исполнений ──
agent = _FakeAgent()
_llm_worker._wire_tool_budget(agent, max_calls=3)
before, after = _budget_hooks(agent)
check("budget: оба хука навешаны (before+after)", before is not None and after is not None)
# 6 вызовов: первые 3 исполняются, остальные отменены
cancels = 0
for _ in range(6):
    ev = _FakeBeforeEvent()
    before(ev)
    if ev.cancel_tool:
        cancels += 1
check("budget: 6 вызовов при лимите 3 → 3 отменены", cancels == 3, str(cancels))
# после партии: end_turn взводится, когда исполнений >= лимита
aev = _FakeAfterEvent()
after(aev)
check("budget: end_turn взведён после пробития лимита", bool(aev.end_turn), str(aev.end_turn))
check("budget: сообщение end_turn инструктивное",
      "бюджет шагов" in str(aev.end_turn) and "оформи задачей" in str(aev.end_turn),
      str(aev.end_turn)[:120])
# после halted: любые новые вызовы отменяются
ev = _FakeBeforeEvent()
before(ev)
check("budget: после halted все вызовы отменены", bool(ev.cancel_tool))

# ── 2. лимит не пробит → end_turn не взводится ──
agent2 = _FakeAgent()
_llm_worker._wire_tool_budget(agent2, max_calls=10)
before2, after2 = _budget_hooks(agent2)
for _ in range(5):
    ev = _FakeBeforeEvent()
    before2(ev)
    check_ok = not ev.cancel_tool
aev2 = _FakeAfterEvent()
after2(aev2)
check("budget: 5 исполнений при лимите 10 → цикл живёт", aev2.end_turn is False,
      str(aev2.end_turn))

# ── 3. лимит 0 — отменяет сразу ──
agent3 = _FakeAgent()
_llm_worker._wire_tool_budget(agent3, max_calls=0)
before3, _ = _budget_hooks(agent3)
ev = _FakeBeforeEvent()
before3(ev)
check("budget: max_calls=0 отменяет первый вызов", bool(ev.cancel_tool))

# ── 4. env-override читается воркером ──
os.environ["SIMBIOT_MAX_TOOL_CALLS"] = "7"
check("budget: env SIMBIOT_MAX_TOOL_CALLS читается",
      int(os.environ.get("SIMBIOT_MAX_TOOL_CALLS") or 30) == 7)
del os.environ["SIMBIOT_MAX_TOOL_CALLS"]
os.environ["SIMBIOT_MAX_TURNS"] = "15"
check("budget: env SIMBIOT_MAX_TURNS читается",
      int(os.environ.get("SIMBIOT_MAX_TURNS") or 15) == 15)
del os.environ["SIMBIOT_MAX_TURNS"]

# ── 5. реальные strands-хуки: create_harness + add_hook не падают ──
try:
    from harness.runtime import setup_provider
    setup_provider()
    from strands_harness import create_harness
    ag = create_harness(
        model=None, instructions="test", session=None,
        builtin_tools=[], memory=False, skills=False,
    )
    ag = _llm_worker._wire_tool_budget(ag, max_calls=30)
    check("budget: хук навешан на реальный strands-агент", True)
except Exception as e:
    check("budget: реальный агент с хуком не падает", False, str(e)[:200])

print()
print("FAILS:", len(fails))
sys.exit(1 if fails else 0)