# 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.

"""Pixel-level Safety Net: VLM comparison of the target crop before execution.

Validates that the visual target button/component at the action coordinates is
still present, visible, and interactive on the live screen compared to the
original (decision-time) screenshot.

Extracted from ``validator.py``; the public entry point remains the
``ValidatorNode`` method, which delegates here. The ``get_llm_fn`` dependency
is injected by the facade so that tests patching
``artemis.agents.validator.validator.get_llm`` keep working.
"""

import asyncio
import base64
import hashlib
import math
from io import BytesIO
from pathlib import Path

from langchain_core.messages import HumanMessage, SystemMessage
from PIL import Image

from artemis.agents.validator.categories import ValidationErrorCategory
from artemis.context import ArtemisContext
from artemis.graph.state import State
from artemis.services.llm import acomplete_structured
from artemis.utils import visualization
from artemis.utils.logger import get_logger

logger = get_logger(__name__)

# Only interactive, element-focused actions are subject to pixel validation.
PIXEL_VALIDATED_ACTIONS = (
    "tap",
    "long_press_on",
    "focus_and_input_text",
    "input_text",
    "click_coordinate",
    "click",
    "long_press",
)

_MAX_ATTEMPTS = 3
_RETRY_DELAY_SECONDS = 0.3


def _annotate_target(image_bytes, coords, coordinate_size, **kwargs) -> bytes:
    """Map device pixels into each encoded screenshot before marking/cropping it."""
    if coordinate_size is not None:
        width, height = coordinate_size
        with Image.open(BytesIO(image_bytes)) as image:
            image_width, image_height = image.size
        if abs((image_width / image_height) / (width / height) - 1) > 0.01:
            raise ValueError("Screenshot orientation/geometry changed; observe the target again.")
        x, y = coords
        if not (0 <= x < width and 0 <= y < height):
            raise ValueError("Target lies outside the observed device screen.")
        coords = [round(x * image_width / width), round(y * image_height / height)]
    return visualization.crop_and_annotate_target(image_bytes, coords, **kwargs)


def _strict_target(action_item: dict) -> bool:
    return any(
        action_item.get(key)
        for key in (
            "target_text",
            "target_resource_id",
            "target_class",
            "target_bounds",
            "target_description",
        )
    ) or action_item.get("action") in {"input_text", "focus_and_input_text"}


def _unverified(action_item: dict, message: str) -> tuple[bool, ValidationErrorCategory, str]:
    if _strict_target(action_item):
        return False, ValidationErrorCategory.PIXEL_VERIFICATION_UNAVAILABLE, message
    return True, ValidationErrorCategory.PIXEL_BYPASSED, ""


def _init_llm_and_prompt(ctx: ArtemisContext, get_llm_fn):
    """Resolves the VLM and loads the prompt template (once per validation)."""
    try:
        llm = get_llm_fn(ctx, name="validator_pixel_safety_net")
    except Exception:
        llm = get_llm_fn(ctx, name="validator")

    prompt_path = Path(__file__).parent.joinpath("pixel_safety_net.md")
    prompt = prompt_path.read_text(encoding="utf-8")
    return llm, prompt


def _describe_target(action_item: dict) -> str:
    """Describe the structured action target for the pixel judge.

    Three kinds, keyed on provenance: an index-addressed target carries observed
    element metadata; a coordinate target carries only the Operator's own
    statement of what it aimed at (``target_description``); a legacy record may
    carry neither.
    """
    label = action_item.get("target_text")
    resource_id = action_item.get("target_resource_id")
    class_name = action_item.get("target_class")
    bounds = action_item.get("target_bounds")
    description = action_item.get("target_description")
    lines = ["[Target]", f"Action: {action_item.get('action')}"]
    if label or resource_id or class_name or bounds:
        lines.append("Kind: specific UI control")
        if label:
            lines.append(f"Label: {label}")
        if resource_id:
            lines.append(f"Resource ID: {resource_id}")
        if class_name:
            lines.append(f"Class: {class_name}")
        if not label and not resource_id:
            lines.append("Label: (unlabelled; identify it by its shape in Image 1)")
        if bounds:
            lines.append(f"Bounds at decision time: {bounds}")
    elif description:
        lines.append("Kind: described target")
        lines.append(f"Operator's description: {description}")
    else:
        lines.append(
            "Kind: coordinates only (no target metadata recorded;"
            " inspect Image 1 for a control at the red dot)"
        )
    return "\n".join(lines)


def _build_messages(
    prompt: str,
    orig_crop_bytes: bytes,
    live_crop_bytes: bytes,
    action_item: dict,
    state: State | None,
) -> list:
    """Formulates the multi-modal verification message pair."""
    orig_b64 = base64.b64encode(orig_crop_bytes).decode("utf-8")
    live_b64 = base64.b64encode(live_crop_bytes).decode("utf-8")
    action_name = action_item.get("action")

    user_content = [
        {"type": "text", "text": _describe_target(action_item)},
        {"type": "text", "text": "[Image 1 (Reference)]"},
        {
            "type": "image_url",
            "image_url": {"url": f"data:image/jpeg;base64,{orig_b64}"},
        },
        {"type": "text", "text": "[Image 2 (Current State)]"},
        {
            "type": "image_url",
            "image_url": {"url": f"data:image/jpeg;base64,{live_b64}"},
        },
    ]
    if state:
        thoughts = []
        native_thought = getattr(state, "operator_native_thinking", None)
        if native_thought and native_thought.strip():
            thoughts.append(native_thought.strip())
        raw_thought = getattr(state, "operator_raw_thinking", None)
        if raw_thought and raw_thought.strip():
            thoughts.append(raw_thought.strip())

        operator_thought = "\n\n".join(thoughts)
        if operator_thought.strip():
            context_text = (
                "[Planned Action & Original Thinking]\nAction:"
                f" {action_name}\nOriginal"
                f" Thinking:\n{operator_thought.strip()}"
            )
            user_content.append({"type": "text", "text": f"\n{context_text.strip()}"})

    return [
        SystemMessage(content=prompt),
        HumanMessage(content=user_content),
    ]


async def _run_attempt(
    session,
    llm,
    prompt: str,
    orig_crop_bytes: bytes,
    current_coords,
    action_item: dict,
    state: State | None,
    attempt: int,
    *,
    coordinate_size: tuple[int, int] | None = None,
) -> tuple[bool, ValidationErrorCategory, str] | None:
    """Runs one pixel validation attempt.

    Returns a final (passed, category, reason) verdict, or None when the
    attempt yielded no verdict (never happens today: every branch is final).
    Raises on infrastructure errors so the caller's retry loop can decide.
    """
    # A. Take fresh live screenshot (captures state changes / dynamic loading settling)
    try:
        live_screenshot_b64 = await session.screenshot_b64()
    except Exception as exc:
        raise RuntimeError(f"Failed to acquire live screenshot: {exc}") from exc

    if not live_screenshot_b64:
        raise RuntimeError("Live screenshot was empty; target identity could not be checked.")

    # B. Decode and crop the live screen target
    live_bytes = base64.b64decode(live_screenshot_b64)
    live_crop_bytes = _annotate_target(
        live_bytes, current_coords, coordinate_size, crop_size=None, dot_radius=15
    )

    # C. Formulate multi-modal verification message
    messages = _build_messages(prompt, orig_crop_bytes, live_crop_bytes, action_item, state)

    # D. Invoke Universal VLM. Parsing (fences, repair) and one
    # corrective re-ask on unparseable JSON are handled by the
    # structured-output layer; a StructuredOutputError lands in
    # the attempt loop's except like any other attempt failure.
    res_json = await acomplete_structured(llm, messages)
    if not isinstance(res_json, dict):
        raise ValueError(
            "Pixel validation response parsed to"
            f" {type(res_json).__name__}, expected a JSON object."
        )
    reasoning = res_json.get("reasoning", "")
    is_present = res_json.get("is_present")
    confidence = res_json.get("confidence")
    if (
        not isinstance(is_present, bool)
        or isinstance(confidence, bool)
        or not isinstance(confidence, (int, float))
        or not math.isfinite(confidence)
        or not 0 <= confidence <= 1
    ):
        raise ValueError(
            "Pixel validation must include boolean is_present and confidence in [0, 1]."
        )
    if _strict_target(action_item) and (
        res_json.get("target_matches") is not True
        or res_json.get("target_kind") not in {"control", "surface"}
    ):
        is_present = False
        reasoning = reasoning or "The intended target was not positively identified."
    if (
        action_item.get("action") in {"input_text", "focus_and_input_text"}
        or any(
            action_item.get(key)
            for key in ("target_text", "target_resource_id", "target_class", "target_bounds")
        )
    ) and res_json.get("target_kind") != "control":
        is_present = False
        reasoning = "A recorded UI control cannot be validated as a blank surface."

    logger.info(
        f"Pixel validation attempt {attempt}/{_MAX_ATTEMPTS} result:"
        f" is_present={is_present}, confidence={confidence:.2f}, reasoning={reasoning}"
    )

    # E. Branching based on outcome
    if is_present and (confidence >= 0.7 or not _strict_target(action_item)):
        logger.info(f"Pixel validation SUCCESS on attempt {attempt}/{_MAX_ATTEMPTS}.")
        return True, ValidationErrorCategory.NONE, ""

    # Target not present: fail or bypass based on confidence without retry!
    if confidence >= 0.7 or _strict_target(action_item):
        reason_msg = (
            f"Pixel-level validation failed: {reasoning}"
            if reasoning
            else (
                "Pixel-level validation failed: The target UI"
                " element at the exact coordinates has disappeared,"
                " changed, or is blocked."
            )
        )
        return (
            False,
            ValidationErrorCategory.PIXEL_TARGET_DISAPPEARED,
            reason_msg,
        )

    logger.warning(
        f"Pixel validation failed but confidence ({confidence:.2f}) is below 0.7. Bypassing check."
    )
    return True, ValidationErrorCategory.PIXEL_BYPASSED, ""


async def validate_action_precondition_pixel(
    ctx: ArtemisContext,
    session,
    action_item: dict,
    pre_screenshot_b64: str,
    original_coords: list[int] | None = None,
    state: State | None = None,
    *,
    get_llm_fn,
) -> tuple[bool, ValidationErrorCategory, str]:
    """Pixel-level Safety Net entry point (see module docstring)."""
    action_name = action_item.get("action")

    # 1. White-list filtering: Only validate interactive, element-focused actions
    if action_name not in PIXEL_VALIDATED_ACTIONS:
        return True, ValidationErrorCategory.PIXEL_BYPASSED, ""

    current_coords = action_item.get("coordinates")
    orig_coords = original_coords or current_coords

    if not orig_coords or len(orig_coords) != 2:
        logger.info("Pixel safety net skipped: no valid coordinates.")
        return _unverified(action_item, "No valid coordinates to verify the intended target.")

    if not current_coords or len(current_coords) != 2:
        current_coords = orig_coords

    device = getattr(ctx, "device", None)
    dimensions = (getattr(device, "device_width", None), getattr(device, "device_height", None))
    coordinate_size = (
        dimensions if all(type(value) is int and value > 0 for value in dimensions) else None
    )

    try:
        neighborhood = _annotate_target(
            base64.b64decode(pre_screenshot_b64),
            orig_coords,
            coordinate_size,
            crop_size=256,
            dot_radius=0,
        )
    except (ValueError, TypeError, OSError):
        neighborhood = pre_screenshot_b64.encode()
    # A changing clock elsewhere on screen is not fresh evidence for a
    # previously rejected blank point. Grounded control metadata still allows
    # a new validation; this does not restrict ordinary repeated actions.
    key = hashlib.sha256(repr(orig_coords).encode() + neighborhood).hexdigest()
    rejected = getattr(ctx, "rejected_pixel_targets", None)
    if not isinstance(rejected, dict):
        rejected = {}
        ctx.rejected_pixel_targets = rejected
    grounded = bool(action_item.get("target_bounds") and action_item.get("target_class"))
    if key in rejected and not grounded:
        return (
            False,
            ValidationErrorCategory.PIXEL_TARGET_DISAPPEARED,
            (
                "This point was already rejected for this observation. A renamed description "
                "is not new evidence; locate the intended control again or obtain fresh target evidence."
            ),
        )

    logger.info(
        "Pre-execution pixel validation: checking target position"
        f" {orig_coords} in original image vs {current_coords} on live"
        " screen..."
    )

    # 2. Decode and crop the static reference (original) image ONCE outside the loop to save CPU
    try:
        orig_bytes = base64.b64decode(pre_screenshot_b64)
        orig_crop_bytes = _annotate_target(
            orig_bytes, orig_coords, coordinate_size, crop_size=None, dot_radius=15
        )
    except Exception as e:
        logger.error(
            f"Failed to prepare original crop for pixel safety net: {e}. Target unverified."
        )
        return _unverified(
            action_item, "Reference screenshot unavailable; target identity is unverified."
        )

    # 3. Get LLM Configuration once outside the loop
    try:
        llm, prompt = _init_llm_and_prompt(ctx, get_llm_fn)
    except Exception as e:
        logger.error(
            f"Failed to initialize LLM/prompt for pixel safety net: {e}. Target unverified."
        )
        return _unverified(
            action_item, "Pixel validator unavailable; target identity is unverified."
        )

    # 4. Start high-precision execution validation retry loop
    last_error = None
    for attempt in range(1, _MAX_ATTEMPTS + 1):
        logger.info(f"Pre-execution pixel validation: Attempt {attempt}/{_MAX_ATTEMPTS}...")
        try:
            verdict = await _run_attempt(
                session,
                llm,
                prompt,
                orig_crop_bytes,
                current_coords,
                action_item,
                state,
                attempt,
                coordinate_size=coordinate_size,
            )
            if verdict is not None:
                if not verdict[0]:
                    rejected[key] = True
                return verdict
        except Exception as attempt_err:
            logger.warning(
                f"Error during pixel validation attempt {attempt}/{_MAX_ATTEMPTS}: {attempt_err}"
            )
            last_error = attempt_err
            if attempt < _MAX_ATTEMPTS:
                await asyncio.sleep(_RETRY_DELAY_SECONDS)
                continue
            logger.error(
                "Pixel validation repeatedly failed with errors. Last"
                f" error: {last_error}. Target identity remains unverified."
            )
            return _unverified(
                action_item,
                f"Pixel verification unavailable after {_MAX_ATTEMPTS} attempts: {last_error}. "
                "The action was not executed; restore observation before locating the target again.",
            )

    return _unverified(action_item, "Pixel validation did not establish target identity.")