"""Bound on the driver's accessibility-tree walk per capture (``computer_use.ax_max_elements``).

``_DEFAULT_MAX_ELEMENTS`` in ``tool.py`` caps the SURFACED element list; the walk that produces it was
unbounded, so a capture paid for every node in the target's tree before trimming. A bound is pure
latency: the bounded tree is a prefix of the unbounded one, so the elements the model sees are
unchanged. Disabled by 0 (driver default), tuned by config, and read through the real loader.
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional

import pytest

from tools.computer_use import cua_backend
from tools.computer_use.cua_backend_capture import _CaptureMixin


class _StubCapture(_CaptureMixin):
    """Capture-lane shell: enough state for ``_gws_args``, no driver, no session."""

    def __init__(self) -> None:
        self._active_pid: Optional[int] = 607
        self._active_window_id: Optional[int] = 382
        self._session_id: Optional[str] = None
        self._last_app = ""

    def _resolve_capture_windows(self, mode: str, app: Optional[str], pid: Optional[int],
                                 window_id: Optional[int]) -> List[Dict[str, Any]]:
        return [{"app_name": "Finder", "pid": 607, "window_id": 382, "title": "", "z_index": 1,
                 "off_screen": False}]

    def _set_active_target(self, target: Dict[str, Any]) -> None:
        self._active_pid, self._active_window_id = target["pid"], target["window_id"]


def _write_config(tmp_path, body: str) -> None:
    """A user config.yaml the backend readers reach through ``load_config`` (not a patched reader)."""
    (tmp_path / "config.yaml").write_text(body, encoding="utf-8")


class TestAxWalkBound:

    def test_configured_value_reaches_the_driver_args(self, tmp_path, monkeypatch):
        """End to end through the loader the backend uses: config.yaml -> get_window_state args."""
        _write_config(tmp_path, "computer_use:\n  ax_max_elements: 350\n")
        monkeypatch.setenv("HERMES_HOME", str(tmp_path))
        assert cua_backend._cua_configured_ax_max_elements() == 350
        stub = _StubCapture()
        assert stub._gws_args()["max_elements"] == 350
        # ... and capture() forwards the bound it actually sent onto the CaptureResult
        monkeypatch.setattr(stub, "_capture_window_state", lambda: (None, None, [], ""))
        assert stub.capture("ax").ax_max_elements == 350

    def test_zero_disables_the_bound(self, tmp_path, monkeypatch):
        """0 restores the driver default and must not leak a ``max_elements`` key into the payload."""
        _write_config(tmp_path, "computer_use:\n  ax_max_elements: 0\n")
        monkeypatch.setenv("HERMES_HOME", str(tmp_path))
        assert cua_backend._cua_configured_ax_max_elements() == 0
        assert "max_elements" not in _StubCapture()._gws_args()


class TestCappedWalkHint:
    """When the driver's walk stopped at the bound, the spill file is not the full tree — say so."""

    @staticmethod
    def _summary(monkeypatch, n_elements: int, bound: int) -> str:
        from tools.computer_use import tool
        from tools.computer_use.backend import CaptureResult, UIElement
        monkeypatch.setattr(tool, "_spill_elements_to_file", lambda cap: "/tmp/elements.json")
        cap = CaptureResult(mode="ax", width=800, height=600, ax_max_elements=bound,
                            elements=[UIElement(index=i + 1, role="AXButton", label=f"b{i}") for i in range(n_elements)])
        return "\n".join(tool._capture_summary_lines(tool._capture_view(cap, max_elements=5)))

    @pytest.mark.parametrize(("n_elements", "bound", "capped"), [(8, 8, True), (8, 50, False), (8, 0, False)])
    def test_hint_names_a_cap_only_when_the_walk_hit_it(self, monkeypatch, n_elements, bound, capped):
        text = self._summary(monkeypatch, n_elements=n_elements, bound=bound)
        assert ("accessibility walk capped" in text) is capped
        assert (f"accessibility walk capped at {bound} elements; pass app= to narrow" in text) is capped
        assert ("full element tree with untruncated labels saved to" in text) is not capped
        assert "element tree with untruncated labels saved to" in text