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

"""Setup phases of ``Explorer.run``.

Split out of ``artemis.agents.explorer.explorer``: the flash-mode fast path
and the named preparation phases executed before the reasoning loop (image
record resolution, version limits, screen dimensions, image pool, GenAI
client, model parameters, prompt construction, and initial screenshot
annotation), packaged as a mixin consumed by ``Explorer``.  Patched
collaborators (``settings``, ``StorageManager``, ``genai``,
``is_ocr_configured``, ``perform_ocr``, ``draw_dots``,
``_run_object_detection``, ``logger``) are resolved through the facade module
at call time; see ``artemis.agents.explorer._facade``.
"""

import base64
import glob
import hashlib
import json
import os
from pathlib import Path
import re

from artemis.agents.explorer._facade import facade
from artemis.agents.explorer.constants import EXPLORE_DESCRIPTIONS
from artemis.agents.explorer.perception_tools import load_detector_templates
from artemis.graph.state import State
from artemis.utils.ocr_xml_fusion import fuse_ocr_with_xml
from artemis.utils.visualization import format_minimal_list_with_points


class RunSetupMixin:
    """Flash mode and pre-loop setup phases of :class:`Explorer`."""

    async def _run_flash(self, query: str, screenshot_path: str) -> str:
        """Flash mode: single-shot object detection without a reasoning loop."""
        _ex = facade()
        templates, global_timeout = load_detector_templates()

        self.screenshot_path = screenshot_path
        self.image_name = None

        queries = [q.strip() for q in query.split("|") if q.strip()]

        try:
            result = await _ex._run_object_detection(
                self.ctx,
                screenshot_path,
                queries,
                templates,
                global_timeout=global_timeout,
            )
            detected_items = result.get("detected", [])
            candidates = []
            for idx, item in enumerate(detected_items):
                pos = item.get("point")
                if pos and isinstance(pos, list) and len(pos) == 2:
                    candidates.append(
                        {
                            "label": f"D{idx + 1}",
                            "coords": pos,
                            "description": item.get("label", query),
                        }
                    )
            fallback_message = "" if candidates else f"Failed to detect: {query}"
            return json.dumps(
                {
                    "candidates": candidates,
                    "fallback_message": fallback_message,
                },
                ensure_ascii=False,
            )
        except Exception as e:
            _ex.logger.error(f"Flash mode object detection failed: {e}")
            return json.dumps(
                {
                    "candidates": [],
                    "fallback_message": f"Flash mode detection error: {e}",
                },
                ensure_ascii=False,
            )

    def _resolve_image_record(self, screenshot_path: str) -> tuple:
        """Computes the screenshot hash and looks it up in the Data Engine DB."""
        _ex = facade()
        image_name = None
        record = None
        try:
            sha256_hash = hashlib.sha256()
            with open(screenshot_path, "rb") as f:
                for byte_block in iter(lambda: f.read(4096), b""):
                    sha256_hash.update(byte_block)
            computed_hash = sha256_hash.hexdigest()
            _ex.logger.info(f"Computed screenshot hash: {computed_hash}")

            # Check if it exists in DB
            db_path = _ex.settings.DATA_ENGINE_DB_PATH
            base_dir = _ex.settings.TRACES_PATH
            storage = _ex.StorageManager(db_path, base_dir)
            record = storage.get_image(computed_hash)

            if record:
                image_name = computed_hash
            else:
                _ex.logger.warning(
                    f"Image hash {computed_hash} not found in Data Engine DB."
                    " Data Engine might not be synced yet."
                )

        except Exception as e:
            _ex.logger.warning(f"Failed to compute hash or check DB: {e}")

        return image_name, record

    def _resolve_version_limits(self, version: str) -> tuple[str, int]:
        """Resolves mode label, iteration budget, and the effective tool denylist."""
        _ex = facade()
        mode = version.capitalize()
        version_denylist = set()
        if version == "ultra":
            max_iterations = 8
        elif version == "pro":
            max_iterations = 3
            version_denylist = {
                "ask_image_processor",
                "get_ocr_list",
                "inspect_region",
                "detect_objects",
            }
        else:
            max_iterations = 3
            version_denylist = {
                "ask_image_processor",
                "get_ocr_list",
                "inspect_region",
            }

        if not _ex.is_ocr_configured():
            version_denylist.add("get_ocr_list")

        try:
            denylisted_config = (
                self.ctx.agent_config.denylisted_tools.get("explorer", [])
                if self.ctx.agent_config
                else []
            )
            self.denylisted_tools = version_denylist.union(set(denylisted_config))
        except (TypeError, AttributeError):
            self.denylisted_tools = version_denylist

        return mode, max_iterations

    def _resolve_screen_dimensions(self, state: State) -> None:
        """Resolves device screen dimensions from operator data or the device."""
        ctx = self.ctx
        self.width = 1080
        self.height = 2400
        operator_raw_data = getattr(state, "operator_raw_data", {}) or {}
        w_raw = operator_raw_data.get("width")
        h_raw = operator_raw_data.get("height")
        if isinstance(w_raw, int) and isinstance(h_raw, int):
            self.width = w_raw
            self.height = h_raw
        else:
            if ctx.device and getattr(ctx.device, "device_width", None):
                self.width = ctx.device.device_width
            if ctx.device and getattr(ctx.device, "device_height", None):
                self.height = ctx.device.device_height

    def _init_image_pool(self, screenshot_path: str) -> None:
        """Seeds the Image Pool with the original screenshot as img_0."""
        self.image_pool = {
            "img_0": {
                "path": screenshot_path,
                "transform": {
                    "offset_x": 0.0,
                    "offset_y": 0.0,
                    "scale_x": 1.0,
                    "scale_y": 1.0,
                },
                "description": "Original complete screenshot",
            }
        }
        self.next_img_id = 1

    def _ensure_genai_client(self):
        """Initializes or reuses the context-level GenAI client for connection pooling."""
        _ex = facade()
        ctx = self.ctx
        client = getattr(ctx, "_genai_client", None)
        if client is None:
            _ex.logger.info(
                "Initializing new GenAI client on context for connection pooling (Explorer)..."
            )
            client = _ex.genai.Client(
                api_key=_ex.settings.GOOGLE_API_KEY.get_secret_value()
                if _ex.settings.GOOGLE_API_KEY
                else None
            )
            ctx._genai_client = client
        return client

    def _resolve_model_params(self) -> tuple:
        """Resolves model name, temperature, thinking level, and fallback model."""
        llm_cfg = self.ctx.llm_config.explorer
        temperature = 0.1
        fallback_model = None
        thinking_level = None
        if llm_cfg:
            model_name = llm_cfg.model
            if "/" in model_name:
                model_name = model_name.split("/")[-1]
            if getattr(llm_cfg, "temperature", None) is not None:
                temperature = llm_cfg.temperature
            if getattr(llm_cfg, "thinking_level", None) is not None:
                thinking_level = llm_cfg.thinking_level
            if getattr(llm_cfg, "fallback", None):
                fallback_model = llm_cfg.fallback.model
                if "/" in fallback_model:
                    fallback_model = fallback_model.split("/")[-1]
        else:
            model_name = "gemini-3.7-flash"

        return model_name, temperature, thinking_level, fallback_model

    def _build_prompt_template(
        self, version: str, mode: str, max_iterations: int
    ) -> tuple[str | None, str | None]:
        """Builds the system prompt template; returns (prompt, error_message)."""
        prompt_path = Path(__file__).parent / "explorer.json"
        if not prompt_path.exists():
            return None, "Error: Explorer prompt template not found."

        try:
            content = prompt_path.read_text(encoding="utf-8")
            content_no_comments = re.sub(r"(?<!:)\/\/.*", "", content)
            content_no_comments = re.sub(r"/\*.*?\*/", "", content_no_comments, flags=re.DOTALL)
            content_no_comments = re.sub(r",\s*([\]}])", r"\1", content_no_comments)
            data = json.loads(content_no_comments)
        except Exception as e:
            return None, f"Error loading or parsing explorer.json: {e}"

        prompt_parts = []
        for section, content_val in data.items():
            prompt_parts.append(f"# {section}")
            if isinstance(content_val, list):
                for bullet in content_val:
                    prompt_parts.append(f"- {bullet}")
            else:
                prompt_parts.append(content_val)
            prompt_parts.append("")

        prompt_template = "\n".join(prompt_parts)
        prompt_template += "\n{denylist}"
        if self.denylisted_tools:
            denylist_section = self.DENYLIST_TEMPLATE.format(
                tools=", ".join(sorted(self.denylisted_tools))
            )
        else:
            denylist_section = ""
        prompt_template = prompt_template.replace("{denylist}", denylist_section)

        explore_info = EXPLORE_DESCRIPTIONS.get(version, EXPLORE_DESCRIPTIONS["pro"])
        version_prompt = explore_info.get("version_prompt")
        if version_prompt:
            version_prompt = version_prompt.format(mode=mode, max_iterations=max_iterations)
            prompt_template += f"\n# EXECUTION CONSTRAINT\n- {version_prompt}\n"

        return prompt_template, None

    async def _build_fused_hierarchy(self, record, state: State, screenshot_path: str):
        """Loads (and if needed OCR-fuses) the UI hierarchy for the screenshot."""
        _ex = facade()
        fused_xml = []
        if record and getattr(record, "ui_tree", None):
            ui_tree = record.ui_tree
            ocr_results = getattr(record, "ocr_result", None)
            if ocr_results is None:
                if _ex.is_ocr_configured():
                    try:
                        _ex.logger.info(
                            "Previous screenshot OCR is missing. Running OCR on-the-fly..."
                        )
                        with open(screenshot_path, "rb") as img_file:
                            img_b64 = base64.b64encode(img_file.read()).decode("utf-8")
                        ocr_results = await _ex.perform_ocr(img_b64, client=self.http_client)
                    except Exception as ocr_err:
                        _ex.logger.error(
                            f"On-the-fly OCR failed for previous screenshot: {ocr_err}"
                        )
                        ocr_results = []
                else:
                    ocr_results = []

            fused_xml = fuse_ocr_with_xml(ui_tree, ocr_results or [])
            _ex.logger.info("Successfully loaded and fused UI hierarchy for previous screenshot.")
        elif hasattr(state, "latest_ui_hierarchy"):
            if screenshot_path == getattr(state, "latest_screenshot", None):
                fused_xml = state.latest_ui_hierarchy
        return fused_xml

    def _annotate_initial_screenshot(
        self, fused_xml, screenshot_path: str, image_name, minimal_list: str
    ) -> tuple[str, object]:
        """Draws labeled dots for the fused hierarchy; returns (minimal_list, marked_path)."""
        _ex = facade()
        ctx = self.ctx
        marked_path = None
        try:
            _ex.logger.info(
                "Explorer self-annotating initial screenshot using latest_ui_hierarchy..."
            )
            formatted_list, points, labels = format_minimal_list_with_points(
                fused_xml, self.width, self.height
            )
            minimal_list = formatted_list
            self.global_label_idx = len(points) + 1

            base_dir = (
                Path(ctx.data_engine.base_dir)
                if ctx.data_engine and getattr(ctx.data_engine, "base_dir", None)
                else None
            )
            if not base_dir:
                db_path = _ex.settings.DATA_ENGINE_DB_PATH
                base_dir = _ex.settings.TRACES_PATH
            images_dir = base_dir / "images"
            initial_marked_dir = images_dir / "initial_marked"
            initial_marked_dir.mkdir(parents=True, exist_ok=True)

            existing_files = glob.glob(
                str(initial_marked_dir / f"{image_name or 'temp_image'}_*.jpg")
            )
            max_seq = 0
            for f in existing_files:
                match = re.search(r"_(\d+)\.jpg$", f)
                if match:
                    max_seq = max(max_seq, int(match.group(1)))
            seq = max_seq + 1
            marked_path = initial_marked_dir / f"{image_name or 'temp_image'}_{seq}.jpg"

            # Draw dots on the raw screenshot
            _ex.draw_dots(screenshot_path, points, labels, str(marked_path))
            _ex.logger.info(
                f"Successfully drew {len(points)} dots and saved marked"
                f" image to {marked_path}"
            )
        except Exception as e:
            _ex.logger.error(f"Failed to self-annotate initial screenshot: {e}")
        return minimal_list, marked_path

    async def _prepare_initial_annotation(
        self, record, state: State, screenshot_path: str, minimal_list: str, image_name
    ) -> tuple[str, str]:
        """Generates initial visual annotations when no minimal list is provided.

        Returns the (possibly updated) minimal list and the image path the
        model should read (marked screenshot when available, raw otherwise).
        """
        marked_path = None
        if not minimal_list:
            fused_xml = await self._build_fused_hierarchy(record, state, screenshot_path)
            if fused_xml:
                minimal_list, marked_path = self._annotate_initial_screenshot(
                    fused_xml, screenshot_path, image_name, minimal_list
                )

        if marked_path and os.path.exists(str(marked_path)):
            img_to_read = str(marked_path)
        else:
            img_to_read = screenshot_path
        return minimal_list, img_to_read