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

import asyncio
import base64
import hashlib
import json
from pathlib import Path

from artemis.context import ArtemisContext
from artemis.controllers.unified_controller import UnifiedMobileController
from artemis.data_engine.trace import trace
from artemis.graph.state import State
from artemis.graph.visibility import strict_state
from artemis.runtime.task_control import TaskControl
from artemis.utils.logger import get_logger
from artemis.utils.ocr_api import is_ocr_configured, perform_ocr
from artemis.utils.ocr_xml_fusion import (
    _crop_image_remove_status_bar,
    _detect_status_bar_height,
    _map_coordinates_back,
    fuse_ocr_with_xml,
)

logger = get_logger(__name__)


def _should_skip_settling(state: State) -> bool:
    """Determine if we should skip screen settling based on the last executed action.

    Purely logical, non-UI-drawing, or wait operations don't require dynamic
    settling checks.
    """
    if not state or not state.structured_decisions:
        # First turn or no prior action, skip settling
        return True

    try:
        actions = json.loads(state.structured_decisions)
        if not actions or not isinstance(actions, list):
            return True

        last_action = actions[-1]
        action_name = last_action.get("action")

        # Actions that are known to have already completed full wait loops or don't render UI
        skip_actions = {
            "wait_for_delay",
        }

        if action_name in skip_actions:
            logger.info(f"Last action was '{action_name}'. Skipping dynamic screen settling.")
            return True

        return False
    except Exception as e:
        logger.warning(
            f"Failed to parse last action for settling heuristic: {e}."
            " Defaulting to settling check."
        )
        return False


@trace(type="agent", name="perception")
def _check_injected_instruction_file(base_dir: str) -> dict | None:
    """Consumes a pending injected-instruction payload, if any.

    Returns the payload dict: `instruction` (str | None) plus the explicit
    `release_loop` stop signal (bool) declared by the injecting caller.
    """
    from artemis.runtime.task_control import consume_instructions

    return consume_instructions(base_dir)


def _unprotect_current_check_lines(ctx: ArtemisContext) -> None:
    """Record the plan's current check lines as editable under user guidance."""
    from artemis.utils.notes import get_note_file_path
    from artemis.utils.plan_grammar import parse_plan

    unprotected = getattr(ctx, "guidance_unprotected_checks", None)
    if not isinstance(unprotected, set):
        unprotected = set()
    try:
        plan_path = get_note_file_path(ctx.data_engine.base_dir, "task_plan")
        text = plan_path.read_text(encoding="utf-8") if plan_path.exists() else ""
    except Exception as e:
        logger.warning(f"Could not read task_plan for the guidance check-line waiver: {e}")
        text = ""
    added = {(ci.kind, ci.text) for ci in parse_plan(text).check_items}
    unprotected |= added
    ctx.guidance_unprotected_checks = unprotected
    logger.info(
        f"User guidance received: {len(added)} declared check line(s) may now be edited"
        " by the Operator."
    )


async def perception_node(state: State, ctx: ArtemisContext) -> dict:
    """Perception node (fast path): captures screenshot, extracts XML hierarchy, and performs OCR.

    Always executes at the start of an operator cycle.
    """
    state = strict_state(state, "perception")
    logger.info("Entering Perception Node...")
    control = getattr(ctx, "task_control", None)
    if isinstance(control, TaskControl):
        await control.checkpoint(
            on_pause=getattr(ctx, "on_task_pause", None),
            on_resume=getattr(ctx, "on_task_resume", None),
        )

    if not ctx.data_engine.current_step_id:
        ctx.data_engine.allocate_step_id()
        logger.info(f"Pre-allocated step ID in Perception Node: {ctx.data_engine.current_step_id}")

    # Check for injected instruction without blocking event loop
    injected_instruction = None
    user_stop_requested = False
    if ctx.data_engine and ctx.data_engine.base_dir:
        injected_payload = await asyncio.to_thread(
            _check_injected_instruction_file, str(ctx.data_engine.base_dir)
        )
        if injected_payload:
            injected_instruction = injected_payload.get("instruction")
            user_stop_requested = bool(injected_payload.get("release_loop"))
            if isinstance(injected_instruction, str) and injected_instruction.strip():
                # User guidance outranks the plan's declared standards: every
                # check line that exists right now loses its machine
                # restoration for the rest of the run (kept on the run context,
                # not the per-turn State), so the Operator can drop or reword
                # it in whichever later write it gets to it. Lines added after
                # this moment stay protected until the next instruction.
                _unprotect_current_check_lines(ctx)

    controller = UnifiedMobileController(ctx)

    # 1. Capture screenshot with hybrid settling strategy
    if _should_skip_settling(state):
        logger.info("Skipping settling. Retrieving screen data directly...")
        device_data = await controller.get_screen_data()
        latest_screenshot_b64 = device_data.base64
        xml_hierarchy = device_data.elements
    else:
        logger.info("Screen settling required. Delaying 0.4s and capturing screen data...")
        await asyncio.sleep(0.4)
        device_data = await controller.get_screen_data()
        latest_screenshot_b64 = device_data.base64
        xml_hierarchy = device_data.elements

    latest_screenshot_bytes = base64.b64decode(latest_screenshot_b64)

    # Update context device dimensions dynamically
    if ctx.device:
        ctx.device.device_width = device_data.width
        ctx.device.device_height = device_data.height
        logger.info(f"Updated context device dimensions: {device_data.width}x{device_data.height}")

    # 3. Perform OCR (if configured, else fallback to pure XML layout)
    ocr_results = []
    if is_ocr_configured():
        screen_height = device_data.height
        status_bar_height = _detect_status_bar_height(xml_hierarchy, screen_height)
        try:
            if status_bar_height > 0:
                logger.info(
                    f"Detected status bar height: {status_bar_height}. Cropping screenshot for OCR."
                )
                cropped_b64, _, _ = await asyncio.to_thread(
                    _crop_image_remove_status_bar,
                    latest_screenshot_b64,
                    status_bar_height,
                )
                raw_ocr_results = await perform_ocr(cropped_b64)
                ocr_results = _map_coordinates_back(raw_ocr_results, status_bar_height)
            else:
                logger.info("No status bar detected. Performing OCR on full screenshot.")
                ocr_results = await perform_ocr(latest_screenshot_b64)
        except Exception as ocr_err:
            logger.warning(
                f"Perception OCR execution encountered error: {ocr_err}. Proceeding with XML-only layout."
            )
            ocr_results = []
    else:
        logger.debug("OCR is not configured. Proceeding with XML-only layout.")

    # 4. Fuse OCR with XML (if ocr_results is empty, returns original xml_hierarchy intact)
    fused_xml = fuse_ocr_with_xml(xml_hierarchy, ocr_results)

    # 5. Save to Data Engine in background thread
    image_name = None
    screenshot_path = None
    if ctx.data_engine:
        image_name = hashlib.sha256(latest_screenshot_bytes).hexdigest()
        screenshot_path = str(ctx.data_engine.get_image_path(image_name))
        asyncio.create_task(
            asyncio.to_thread(
                ctx.data_engine.get_or_create_image,
                latest_screenshot_bytes,
                ui_tree=xml_hierarchy,
                ocr_result=ocr_results,
            )
        )
        logger.info(f"Offloaded perception data saving to background with image_name: {image_name}")

    # Return state update
    update = {
        "latest_screenshot": screenshot_path,
        "latest_ui_hierarchy": fused_xml,
        "operator_raw_data": {
            "screenshot_b64": latest_screenshot_b64,
            "xml_hierarchy": xml_hierarchy if xml_hierarchy is not None else [],
            "ocr_results": ocr_results,
            "width": device_data.width,
            "height": device_data.height,
        },
        "injected_instruction": injected_instruction,
        # sticky_or reducer: once latched True, later False updates keep it True
        "user_stop_requested": user_stop_requested,
    }
    return update