import json
import logging
import platform
import re

from mm_agents.vlaa_gui.core.module import BaseModule
from mm_agents.vlaa_gui.memory.procedural_memory import PROCEDURAL_MEMORY
from mm_agents.vlaa_gui.utils.common_utils import call_llm_safe

logger = logging.getLogger("desktopenv.agent")


class GateAgent(BaseModule):
    """Proposes UI-observable success criteria for a task at step 1.

    Called once at the start of a task to produce 1-3 success criteria that the
    worker then checks every step.  This keeps the large criteria-writing
    instructions out of the worker system prompt.
    """

    def __init__(self, engine_params: dict, platform: str = platform.system().lower()) -> None:
        super().__init__(engine_params, platform)
        self.agent = self._create_agent(
            PROCEDURAL_MEMORY.construct_gate_prompt(platform)
        )

    # ------------------------------------------------------------------
    # JSON extraction helpers (same pattern as VerifierAgent)
    # ------------------------------------------------------------------

    @staticmethod
    def _extract_first_json_dict(text: str) -> dict | None:
        if not text:
            return None
        candidate = text.strip()
        try:
            data = json.loads(candidate)
            if isinstance(data, dict):
                return data
        except Exception:
            pass
        starts = [m.start() for m in re.finditer(r"\{", candidate)]
        for start in starts:
            depth = 0
            in_string = False
            escape = False
            for end in range(start, len(candidate)):
                char = candidate[end]
                if in_string:
                    if escape:
                        escape = False
                    elif char == "\\":
                        escape = True
                    elif char == '"':
                        in_string = False
                    continue
                if char == '"':
                    in_string = True
                    continue
                if char == "{":
                    depth += 1
                    continue
                if char == "}":
                    depth -= 1
                    if depth == 0:
                        snippet = candidate[start : end + 1]
                        try:
                            data = json.loads(snippet)
                            if isinstance(data, dict):
                                return data
                        except Exception:
                            pass
                        break
        return None

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def propose_criteria(self, instruction: str, observation: dict) -> list[dict]:
        """Propose success criteria for a task.

        Args:
            instruction: The task instruction.
            observation: Current UI observation (must contain 'screenshot').

        Returns:
            List of criterion dicts, each with 'text' (str) and
            'exact_check' (bool).  Falls back to a generic criterion on
            parse failure.
        """
        self.agent.reset()
        prompt = (
            f"Task: {instruction}\nPropose success criteria based on the screenshot."
        )
        self.agent.add_message(
            text_content=prompt,
            image_content=observation.get("screenshot"),
            role="user",
        )
        response = call_llm_safe(self.agent, temperature=0.2, agent_type="gate")

        criteria = self._parse_criteria(response)
        return criteria

    def _parse_criteria(self, response: str) -> list[dict]:
        """Parse criteria from LLM response, with fallback."""
        data = self._extract_first_json_dict(response)
        if data and "criteria" in data and isinstance(data["criteria"], list):
            result = []
            for item in data["criteria"]:
                if isinstance(item, dict) and "text" in item:
                    result.append(
                        {
                            "text": str(item["text"]),
                            "exact_check": bool(item.get("exact_check", False)),
                        }
                    )
            if result:
                return result[:3]

        # Fallback: extract bullet points from raw text
        fallback = []
        for line in (response or "").split("\n"):
            line = line.strip().lstrip("-•*0123456789.) ")
            if line and len(line) > 10:
                exact = (
                    "[EXACT CHECK]" in line.upper() or "[EXACT_CHECK]" in line.upper()
                )
                fallback.append({"text": line, "exact_check": exact})
        if fallback:
            return fallback[:3]

        # Last resort
        return [
            {
                "text": "Task requirements visibly satisfied on screen",
                "exact_check": False,
            }
        ]

    @staticmethod
    def format_criteria_for_worker(criteria: list[dict]) -> str:
        """Format criteria into a compact string for injection into worker context.

        Returns a block like:
            SUCCESS CRITERIA (proposed by gate agent — check each step):
            1. [EXACT CHECK] The font size field shows "14pt"
            2. The document title bar shows the new filename
        """
        if not criteria:
            return ""
        lines = ["SUCCESS CRITERIA (proposed by gate agent — check each step):"]
        for i, c in enumerate(criteria, 1):
            tag = "[EXACT CHECK] " if c.get("exact_check") else ""
            lines.append(f"  {i}. {tag}{c['text']}")
        return "\n".join(lines)
