# Copyright 2026 Google LLC
#
# 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.

"""Tests for pixel safety-net target descriptions."""

from pathlib import Path
from io import BytesIO
from unittest.mock import MagicMock
from unittest.mock import AsyncMock
import base64
import pytest
from PIL import Image

from artemis.agents.validator import precondition_pixel as pp
from langchain_core.messages import HumanMessage, SystemMessage

_PNG = b"fake-image-bytes"


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "response",
    [
        {"is_present": True, "confidence": 0.99, "target_kind": "surface", "target_matches": True},
        {"is_present": True, "confidence": 0.99},
        {"is_present": True, "confidence": 0.2, "target_kind": "control", "target_matches": True},
    ],
)
@pytest.mark.parametrize(
    "target",
    [
        {"action": "tap", "target_text": "Input"},
        {"action": "input_text", "target_description": "Input"},
    ],
)
async def test_named_input_is_never_approved_as_blank_or_unverified(response, target, monkeypatch):
    monkeypatch.setattr(pp.visualization, "crop_and_annotate_target", lambda *a, **k: _PNG)
    monkeypatch.setattr(pp, "acomplete_structured", AsyncMock(return_value=response))
    session = AsyncMock()
    session.screenshot_b64.return_value = base64.b64encode(_PNG).decode()
    verdict = await pp._run_attempt(session, None, "prompt", _PNG, [1, 2], target, None, 1)
    assert verdict[0] is False


@pytest.mark.asyncio
async def test_renaming_rejected_point_cannot_dispatch_again(monkeypatch):
    monkeypatch.setattr(pp, "_init_llm_and_prompt", lambda *a: (None, "prompt"))
    attempt = AsyncMock(
        return_value=(False, pp.ValidationErrorCategory.PIXEL_TARGET_DISAPPEARED, "No input")
    )
    monkeypatch.setattr(pp, "_run_attempt", attempt)
    ctx = MagicMock()
    image = Image.new("RGB", (400, 400), "white")
    frames = []
    for color in ("black", "blue"):
        image.putpixel((5, 5), Image.new("RGB", (1, 1), color).getpixel((0, 0)))
        encoded = BytesIO()
        image.save(encoded, format="PNG")
        frames.append(base64.b64encode(encoded.getvalue()).decode())
    for description, frame in zip(("pairing input", "blank surface"), frames, strict=True):
        verdict = await pp.validate_action_precondition_pixel(
            ctx,
            AsyncMock(),
            {"action": "tap", "coordinates": [200, 300], "target_description": description},
            frame,
            get_llm_fn=lambda: None,
        )
        assert verdict[0] is False
    assert attempt.await_count == 1


def _texts(messages: list) -> list[str]:
    human = messages[1]
    return [part["text"] for part in human.content if part["type"] == "text"]


def test_named_control_target_block_precedes_images():
    item = {
        "action": "tap",
        "coordinates": [1001, 718],
        "target_text": "进入全屏模式",
        "target_resource_id": "com.google.android.youtube:id/fullscreen_button",
        "target_class": "android.widget.ImageView",
        "target_bounds": [922, 655, 1080, 781],
    }
    state = MagicMock()
    state.operator_native_thinking = "identified 进入全屏模式"
    state.operator_raw_thinking = None

    messages = pp._build_messages("SYSTEM", _PNG, _PNG, item, state)

    assert isinstance(messages[0], SystemMessage)
    assert isinstance(messages[1], HumanMessage)
    first_text = messages[1].content[0]["text"]
    assert first_text.startswith("[Target]")
    assert "Kind: specific UI control" in first_text
    assert "Label: 进入全屏模式" in first_text
    assert "Resource ID: com.google.android.youtube:id/fullscreen_button" in first_text
    assert "Class: android.widget.ImageView" in first_text
    assert "[922, 655, 1080, 781]" in first_text
    assert messages[1].content[1]["text"] == "[Image 1 (Reference)]"
    assert any("Original Thinking" in t for t in _texts(messages))
    assert any("Action: tap" in t for t in _texts(messages))


def test_coordinates_only_target_is_labelled_as_surface():
    item = {"action": "tap", "coordinates": [540, 400]}

    messages = pp._build_messages("SYSTEM", _PNG, _PNG, item, state=None)

    first_text = messages[1].content[0]["text"]
    assert "Kind: coordinates only" in first_text
    # No hit test runs any more: the block states that nothing was recorded, it
    # does not claim the point was looked up in the UI hierarchy.
    assert "no target metadata recorded" in first_text
    assert "UI hierarchy" not in first_text
    assert "Label:" not in first_text
    assert "Resource ID:" not in first_text
    assert not any("Original Thinking" in t for t in _texts(messages))


def test_label_alone_counts_as_named_control():
    item = {"action": "tap", "coordinates": [10, 10], "target_text": "Submit"}
    text = pp._describe_target(item)
    assert "Kind: specific UI control" in text
    assert "Label: Submit" in text
    assert "Resource ID" not in text


def test_unlabelled_indexed_element_is_still_a_control():
    item = {
        "action": "tap",
        "coordinates": [1001, 718],
        "target_class": "android.widget.ImageView",
        "target_bounds": [922, 655, 1080, 781],
    }
    text = pp._describe_target(item)
    assert "Kind: specific UI control" in text
    assert "unlabelled" in text
    assert "Class: android.widget.ImageView" in text
    assert "[922, 655, 1080, 781]" in text


def test_described_coordinate_target_is_its_own_kind():
    """The model's description is presented as its belief, not as observed metadata."""
    item = {"action": "tap", "coordinates": [540, 400], "target_description": "play button"}
    text = pp._describe_target(item)
    assert "Kind: described target" in text
    assert "Operator's description: play button" in text
    assert "Kind: specific UI control" not in text
    assert "Label:" not in text


def test_observed_metadata_outranks_description():
    item = {
        "action": "tap",
        "coordinates": [10, 10],
        "target_text": "Submit",
        "target_description": "ignored",
    }
    text = pp._describe_target(item)
    assert "Kind: specific UI control" in text
    assert "Label: Submit" in text
    assert "described target" not in text


def test_prompt_rules_are_keyed_on_target_kind():
    prompt = Path(pp.__file__).with_name("pixel_safety_net.md").read_text(encoding="utf-8")
    assert "`Kind: specific UI control`" in prompt
    assert "`Kind: described target`" in prompt
    assert "`Kind: coordinates only`" in prompt
    assert "hidden" in prompt
    assert "Identify what appears under the red dot" in prompt