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

"""Asynchronous Visual Step Summarizer for Artemis Flash profile.

Processes Before/After screenshot pairs and action metadata in the background
to generate high-density, strictly objective visual state transition summaries.

Scheduling (zero-blocking dispatch, bounded retry, bounded flush, step_id
keying with tool_call_id aliases) lives in the shared
:class:`artemis.memory.step_memory.StepMemoryService`; this class contributes
only the visual-transition lens: red action-marker overlay on the BEFORE
frame, the ``flash_summarizer.md`` neutral-wording contract, and versioned
``summary_status`` writes to the DataEngine.
"""

import asyncio
import base64
from pathlib import Path
from typing import Any
from uuid import UUID

from jinja2 import Template
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage

from artemis.context import ArtemisContext
from artemis.memory.step_memory import JobKey, StepMemoryService
from artemis.services.llm import RobustChatModelWrapper, get_google_llm, get_llm
from artemis.services.token_meter import record_llm_usage
from artemis.utils.logger import get_logger
from artemis.utils.visualization import draw_action_overlay_on_image

logger = get_logger(__name__)

# Degenerate-output guard (§5 echo validation): a summary that echoes the
# input's section markers, or falls outside sane length bounds, fails the
# attempt so the bounded service retry regenerates it.
SUMMARY_ECHO_MARKER = "--- ["
SUMMARY_MIN_CHARS = 15
SUMMARY_MAX_CHARS = 1500


def degenerate_summary_reason(text: str) -> str | None:
    """Why a lens output is unusable (echo/length), or None when it is fine."""
    if SUMMARY_ECHO_MARKER in text or text.startswith("---"):
        return "echoes an input section marker"
    if len(text) < SUMMARY_MIN_CHARS:
        return f"too short ({len(text)} chars < {SUMMARY_MIN_CHARS})"
    if len(text) > SUMMARY_MAX_CHARS:
        return f"too long ({len(text)} chars > {SUMMARY_MAX_CHARS})"
    return None


class VisualStepSummarizer(StepMemoryService):
    """Visual-transition lens on top of the shared step-memory runtime.

    Key design properties (inherited from StepMemoryService):
    1. Zero-blocking dispatch: The main Flash runner dispatches and proceeds immediately.
    2. Lossless pending state: The compressor retains the source image until its summary is ready.
    3. Independent bounded retry: Every action owns a retry loop capped at 1 + retry_limit attempts.
    """

    def __init__(
        self,
        ctx: ArtemisContext,
        model_name: str | None = None,
        retry_limit: int = 3,
        *,
        max_concurrency: int = 1,
        flush_timeout_s: float = 30.0,
    ):
        super().__init__(
            ctx,
            max_concurrency=max_concurrency,
            retry_limit=retry_limit,
            flush_timeout_s=flush_timeout_s,
        )

        # Initialize lightweight VLM: prioritize explicit model_name
        target_model = model_name or "gemini-2.5-flash-lite"
        self._model_name = target_model
        try:
            if model_name:
                self._llm = get_google_llm(model_name=target_model, temperature=0.0)
            else:
                self._llm = get_llm(ctx, name="summarizer", is_utils=True)
        except Exception:
            self._llm = get_google_llm(model_name=target_model, temperature=0.0)
        try:
            configured = getattr(self._llm, "model", None) or getattr(
                self._llm, "model_name", None
            )
            if isinstance(configured, str) and configured:
                self._model_name = configured
        except Exception:
            pass

        # Load system prompt templates: the dual-frame transition prompt and
        # the single-frame variant (§5 revision: describe whichever frames
        # exist — a Pro step record never carries an independent after-frame,
        # by design, and must not be prompted as if one were expected).
        prompt_path = Path(__file__).parent / "flash_summarizer.md"
        if prompt_path.exists():
            self._prompt_template = prompt_path.read_text(encoding="utf-8")
        else:
            self._prompt_template = (
                "You are the Step Summarizer for an Android UI automation agent.\n"
                "Synthesize the physical action and visual delta between BEFORE and AFTER screens in exactly ONE "
                "continuous first-person paragraph using 'I' (e.g., 'In Step {{ step_number }}, I tapped... and observed...').\n"
                "Strictly avoid subjective validation words: successfully, completed, failed, achieved, navigated to."
            )
        single_path = Path(__file__).parent / "flash_summarizer_single.md"
        if single_path.exists():
            self._single_prompt_template = single_path.read_text(encoding="utf-8")
        else:
            self._single_prompt_template = (
                "You are the Step Summarizer for an Android UI automation agent.\n"
                "Exactly ONE screenshot (the decision frame) is available; there is NO after-action screenshot —"
                " that only means no independent post-action evidence exists.\n"
                "Describe strictly what THIS screen shows and where the action landed (red marker), in exactly ONE"
                " continuous first-person paragraph using 'I' for Step {{ step_number }}. Never describe or guess"
                " the post-action state.\n"
                "Strictly avoid subjective validation words: successfully, completed, failed, achieved, navigated to."
            )

    def dispatch(
        self,
        step_number: int,
        action_name: str,
        action_args: dict[str, Any],
        pre_img_bytes: bytes | None,
        post_img_bytes: bytes | None,
        exec_outcome: str,
        *,
        action_key: str | None = None,
        data_engine_step_id: UUID | str | None = None,
    ) -> None:
        """Dispatches an asynchronous summarization task without blocking the caller.

        Jobs are keyed by the DataEngine step id when one is available; the
        tool_call_id (``action_key``) is retained as an alias so the message
        compressor can keep querying by it. Callers with neither provide the
        step ordinal, matching the legacy keying.
        """
        key: JobKey
        aliases: tuple[JobKey, ...] = ()
        if data_engine_step_id is not None:
            key = str(data_engine_step_id)
            if action_key is not None:
                aliases = (action_key,)
        else:
            key = action_key if action_key is not None else step_number

        payload = {
            "step_number": step_number,
            "action_name": action_name,
            "action_args": action_args,
            "pre_img_bytes": pre_img_bytes,
            "post_img_bytes": post_img_bytes,
            "exec_outcome": exec_outcome,
            "data_engine_step_id": data_engine_step_id,
        }
        self.submit(key, payload, aliases=aliases)

    def _meter_lens_call(self, response: Any) -> None:
        """Meter one raw-model lens call as an ``llm_usage`` trace, best-effort.

        Gateway-wrapped models already meter at the wrapper exit; only the raw
        ``get_google_llm`` bypass needs explicit metering here. Lens prompts
        are tiny and must not overwrite the session's ``last_prompt_tokens``
        (the compaction thresholds' live context base), hence
        ``update_last_prompt=False``.
        """
        if isinstance(self._llm, RobustChatModelWrapper):
            return
        engine = getattr(self.ctx, "data_engine", None) if self.ctx else None
        record_llm_usage(
            engine,
            response,
            source=f"lens:visual_transition:{self._model_name}",
            update_last_prompt=False,
        )

    def _on_status(self, key: JobKey, status: str) -> None:
        """Best-effort DataEngine summary-status write (pending/failed)."""
        if not self.ctx.data_engine:
            return
        payload = self._step_inputs.get(key) or {}
        target_step = payload.get("data_engine_step_id")
        if not target_step:
            return
        try:
            self.ctx.data_engine.update_step_summary(
                target_step,
                None,
                status=status,
                source="visual_transition",
                model=self._model_name,
            )
        except Exception as de_err:
            logger.debug(f"DataEngine summary status update skipped: {de_err}")

    async def _attempt(self, key: JobKey) -> bool:
        """Execute one lightweight VLM attempt for a step transition."""
        input_data = self._step_inputs.get(key)
        if not input_data:
            return False

        step_number = input_data["step_number"]
        action_name = input_data["action_name"]
        action_args = input_data["action_args"]
        pre_bytes = input_data["pre_img_bytes"]
        post_bytes = input_data["post_img_bytes"]
        exec_outcome = input_data["exec_outcome"]

        try:
            # §5 revision: the lens input adapts to whichever frames exist.
            # Both frames -> the classic BEFORE/AFTER transition prompt; a
            # single frame (Pro steps never carry an independent after-frame,
            # by design) -> the single-frame variant, which has no AFTER
            # ACTION section and forbids synthesizing a transition.
            dual = bool(pre_bytes) and bool(post_bytes)
            template = self._prompt_template if dual else self._single_prompt_template
            rendered_prompt = Template(template).render(step_number=step_number)
            content_blocks: list[dict[str, Any]] = [
                {
                    "type": "text",
                    "text": (
                        f"Step {step_number} Physical Action: {action_name}({action_args})\n"
                        f"Controller Outcome: {exec_outcome}"
                    ),
                }
            ]

            if pre_bytes:
                # 🎨 Visually mark the exact action (tap ripple, sequence numbers, swipe arrow) on the BEFORE screenshot
                annotated_pre_bytes = draw_action_overlay_on_image(
                    image_bytes=pre_bytes,
                    action_name=action_name,
                    action_args=action_args,
                )
                b64_pre = base64.b64encode(annotated_pre_bytes).decode("utf-8")
                content_blocks.append(
                    {
                        "type": "text",
                        "text": (
                            "--- [1] BEFORE ACTION SCREEN (Action Marked Visually in Red) ---"
                            if dual
                            else "--- [1] DECISION FRAME: SCREEN AT ACTION TIME"
                            " (Action Marked Visually in Red) ---"
                        ),
                    }
                )
                content_blocks.append(
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64_pre}"}}
                )

            if post_bytes:
                b64_post = base64.b64encode(post_bytes).decode("utf-8")
                content_blocks.append(
                    {
                        "type": "text",
                        "text": (
                            "--- [2] AFTER ACTION SCREEN ---"
                            if dual
                            else "--- [1] SCREEN OBSERVED AFTER THE ACTION"
                            " (no decision frame available) ---"
                        ),
                    }
                )
                content_blocks.append(
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{b64_post}"},
                    }
                )

            messages: list[BaseMessage] = [
                SystemMessage(content=rendered_prompt),
                HumanMessage(content=content_blocks),
            ]

            response = await asyncio.wait_for(self._llm.ainvoke(messages), timeout=25.0)
            self._meter_lens_call(response)
            summary_raw = response.content if isinstance(response.content, str) else ""
            if isinstance(response.content, list):
                summary_raw = "".join(
                    b.get("text", "")
                    for b in response.content
                    if isinstance(b, dict) and "text" in b
                )

            summary_text = summary_raw.strip()
            degenerate = degenerate_summary_reason(summary_text) if summary_text else None
            if degenerate:
                logger.warning(
                    f"VisualStepSummarizer: Discarding degenerate summary for"
                    f" Step {step_number} ({degenerate}): {summary_text[:80]!r}"
                )
                return False
            if summary_text:
                self._summaries[key] = summary_text
                logger.info(
                    f"VisualStepSummarizer: Generated summary for Step {step_number}: {summary_text[:80]}..."
                )

                # Free binary image buffers from memory once summary is secured
                if key in self._step_inputs:
                    self._step_inputs[key]["pre_img_bytes"] = None
                    self._step_inputs[key]["post_img_bytes"] = None

                # Update DataEngine telemetry if active
                if self.ctx.data_engine:
                    try:
                        target_step = input_data.get("data_engine_step_id") or step_number
                        self.ctx.data_engine.update_step_summary(
                            target_step,
                            summary_text,
                            status="ready",
                            source="visual_transition",
                            model=self._model_name,
                        )
                    except Exception as de_err:
                        logger.debug(f"DataEngine step summary update skipped: {de_err}")
                return True

        except Exception as e:
            logger.warning(
                f"VisualStepSummarizer: Error generating summary for step {step_number}: {e}"
            )
        return False