import json
import logging
import platform
import re
import textwrap

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

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


class VerifierAgent(BaseModule):
    """A lightweight judge that checks if the task is truly completed."""

    _UNCERTAINTY_PHRASES = (
        "not sure",
        "unsure",
        "unclear",
        "cannot verify",
        "can't verify",
        "unable to verify",
        "cannot tell",
        "can't tell",
        "hard to tell",
        "not enough information",
        "insufficient information",
        "no clear evidence",
    )

    # -- Task category definitions ----------------------------------------
    _TASK_CATEGORIES: dict[str, dict] = {
        "file_save_office": {
            "keywords": [
                ["libreoffice"],
                ["presentation"],
                ["slide"],
                ["impress"],
                ["spreadsheet"],
                ["libreoffice calc"],
                ["libreoffice writer"],
                [".pptx"],
                [".docx"],
                [".xlsx"],
                [".ods"],
                [".odt"],
                [".odp"],
                [".csv"],
            ],
            "hint": (
                "FILE SAVE CHECK: This task involves a document/spreadsheet/"
                "presentation. The file will be checked programmatically by "
                "comparing its contents against a reference — visual appearance "
                "in the editor is NOT sufficient. Verify that the file was "
                "explicitly SAVED (Ctrl+S, File > Save, or via code). If there "
                "is no evidence of a save action in the trajectory, mark INCOMPLETE."
            ),
        },
        "image_export": {
            "keywords": [
                ["gimp"],
                ["image editor"],
                ["export"],
                [".xcf"],
                [".png", "save"],
                [".jpg", "save"],
                [".jpeg", "save"],
                [".bmp", "save"],
                [".tiff", "save"],
                [".svg", "save"],
            ],
            "hint": (
                "IMAGE EXPORT CHECK: This task involves image editing. The saved "
                "image file will be compared by structural similarity or pixel "
                "analysis — visual appearance in the editor is NOT checked. The "
                "image MUST be exported/saved to disk (File > Export As or "
                "File > Overwrite). If there is no evidence of an export/save "
                "in the trajectory, mark INCOMPLETE."
            ),
        },
        "browser_state": {
            "keywords": [
                ["chrome"],
                ["browser"],
                ["bookmark"],
                ["browser", "tab"],
                ["chrome", "tab"],
                ["open", "tab"],
                ["close", "tab"],
                ["new tab"],
                ["cookie"],
                ["browsing history"],
                ["browser history"],
                ["clear history"],
                ["delete history"],
                ["homepage"],
                ["search engine"],
                ["new tab page"],
            ],
            "hint": (
                "BROWSER STATE CHECK: Browser state (URLs, tabs, bookmarks, "
                "extensions, cookies, history, settings) is checked "
                "programmatically — not by looking at the screenshot. The "
                "evaluator reads actual browser state via accessibility tree, "
                "CLI, or config files. Visual confirmation of 'it looks right' "
                "is NOT sufficient. Verify from the trajectory that the actual "
                "state change was made."
            ),
        },
        "vscode_settings": {
            "keywords": [
                ["vs code"],
                ["vscode"],
                ["visual studio code"],
                ["settings.json"],
                ["keybindings"],
            ],
            "hint": (
                "VS CODE CHECK: VS Code settings, keybindings, and extensions "
                "are checked by reading JSON config files or running CLI commands "
                "(e.g., 'code --list-extensions'). Changes must be persisted to "
                "disk — just opening a settings UI is NOT enough. Verify the "
                "trajectory shows the setting was actually written/saved."
            ),
        },
        "thunderbird_config": {
            "keywords": [
                ["thunderbird"],
                ["email client"],
                ["mail client"],
            ],
            "hint": (
                "THUNDERBIRD CHECK: Thunderbird preferences (prefs.js), message "
                "filters, and folder structure are checked by reading config files "
                "on disk. Changes MUST be saved — the application should be closed "
                "or the setting committed so prefs.js is written. Visual "
                "confirmation in the UI is NOT sufficient."
            ),
        },
        "vlc_config": {
            "keywords": [
                ["vlc"],
                ["media player"],
                ["vlc media"],
            ],
            "hint": (
                "VLC CHECK: VLC settings and state are checked by reading config "
                "files or querying application state programmatically. Settings "
                "must be saved/applied — just opening a preferences dialog is NOT "
                "enough."
            ),
        },
        "system_config": {
            "keywords": [
                ["gsettings"],
                ["dconf"],
                ["timezone"],
                ["locale"],
                ["wallpaper"],
                ["screen resolution"],
                ["night light"],
                ["dark mode"],
                ["system settings"],
                ["notification"],
                ["favorite apps"],
                ["default application"],
            ],
            "hint": (
                "SYSTEM CONFIG CHECK: OS-level settings are checked by running "
                "shell commands (gsettings, dconf, etc.) or reading config files. "
                "Changes must be actually applied at the system level — just "
                "navigating to a settings panel is NOT enough. Verify the "
                "trajectory shows the setting was changed via command or GUI "
                "with a confirmation."
            ),
        },
        "file_operation": {
            "keywords": [
                ["create", "file"],
                ["create", "folder"],
                ["create", "directory"],
                ["move", "file"],
                ["copy", "file"],
                ["rename", "file"],
                ["delete", "file"],
                ["extract", "archive"],
                ["compress"],
                [".zip"],
                [".tar"],
            ],
            "hint": (
                "FILE OPERATION CHECK: File system operations (create, move, copy, "
                "rename, delete) are verified by checking actual file existence and "
                "contents on disk. Visual confirmation in a file manager is NOT "
                "sufficient — the evaluator runs shell commands to check."
            ),
        },
        "code_task": {
            "keywords": [
                ["write", "code"],
                ["write", "script"],
                ["write", "program"],
                ["python", "file"],
                ["bash", "script"],
                [".py"],
                ["test suite"],
                ["unit test"],
            ],
            "hint": (
                "CODE CHECK: The code file will be tested programmatically "
                "(e.g., a test suite will be run against it). The file must be "
                "saved and syntactically correct. Visual appearance in the editor "
                "is NOT checked."
            ),
        },
    }

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

    # -- Task categorization ----------------------------------------------

    @classmethod
    def _categorize_task(cls, instruction: str, trajectory: str = "") -> list:
        """Classify a task into zero or more categories based on instruction text.

        Uses word-boundary matching to avoid false positives from substrings
        (e.g., "calc" should not match "calculates").

        Returns a list of matching category names (may be empty).
        """
        text = (instruction + " " + trajectory).lower()
        matched = []
        for category, spec in cls._TASK_CATEGORIES.items():
            for keyword_group in spec["keywords"]:
                if all(
                    re.search(r"\b" + re.escape(kw) + r"\b", text)
                    if " " not in kw
                    else kw in text
                    for kw in keyword_group
                ):
                    matched.append(category)
                    break  # one matching group is enough for this category
        return matched

    _INFEASIBILITY_SIGNALS = (
        "not possible",
        "impossible",
        "infeasible",
        "cannot be done",
        "can't be done",
        "no such option",
        "no such feature",
        "does not exist",
        "doesn't exist",
        "not supported",
        "not available",
        "no way to",
        "option not found",
        "feature not found",
        "could not find",
        "couldn't find",
        "this cannot be accomplished",
        "unable to complete",
        "unable to accomplish",
        "there is no",
        # Workaround signals — agent may be bypassing infeasibility
        "virtual adapter",
        "virtual device",
        "virtual bluetooth",
        "no battery",
        "no bluetooth",
        "no hardware",
        "power_supply/ is empty",
        "power_supply is empty",
        "installed extension",
        "install extension",
        "pip install",
        "undefined variable",
        "variable is not defined",
        "variable is not set",
        "variable is empty",
    )

    _INFEASIBILITY_HINT = (
        "INFEASIBILITY CHECK: Some desktop tasks are infeasible "
        "(the requested feature/option does not exist, or the action is "
        "impossible in the application). The trajectory contains signals that "
        "the agent may have determined this task is infeasible. If the agent "
        "explored the UI thoroughly and the requested capability genuinely "
        "does not exist, the correct action is fail(), NOT done(). If the "
        "agent called done() but the trajectory shows it could not actually "
        "accomplish the request, mark INCOMPLETE."
    )

    @classmethod
    def _check_infeasibility_signals(cls, trajectory: str) -> bool:
        """Check if the trajectory contains signals the task may be infeasible."""
        traj_lower = trajectory.lower()
        return any(sig in traj_lower for sig in cls._INFEASIBILITY_SIGNALS)

    @classmethod
    def _build_category_hints(
        cls, categories: list, has_infeasibility_signals: bool = False
    ) -> str:
        """Build a combined hint string from matched categories."""
        hints = []
        if has_infeasibility_signals:
            hints.append(cls._INFEASIBILITY_HINT)
        for cat in categories:
            spec = cls._TASK_CATEGORIES.get(cat)
            if spec:
                hints.append(spec["hint"])
        if not hints:
            return ""
        return "\n".join(f"- {h}" for h in hints)

    def _to_bool(self, value) -> bool:
        if isinstance(value, bool):
            return value
        if isinstance(value, str):
            return value.strip().lower() in {
                "true",
                "yes",
                "y",
                "complete",
                "completed",
                "done",
                "1",
            }
        return bool(value)

    def _extract_first_json_dict(self, text: str) -> dict | None:
        if not text:
            return None

        candidate = text.strip()

        # Fast-path: already JSON.
        try:
            data = json.loads(candidate)
            if isinstance(data, dict):
                return data
        except Exception:
            pass

        # Common case: JSON wrapped in markdown/code fences or extra prose.
        # Scan for the first balanced {...} substring that parses as a dict.
        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

    def _normalize_missing_steps(self, missing_steps: str) -> str:
        cleaned = (missing_steps or "").strip()
        lowered = cleaned.lower()
        if lowered in {"", "none", "n/a", "na", "null"}:
            return ""
        if lowered.startswith("none"):
            return ""
        if "no missing" in lowered:
            return ""
        if "nothing" == lowered:
            return ""
        return cleaned

    def _contains_uncertainty(self, text: str) -> bool:
        lowered = (text or "").lower()
        return any(phrase in lowered for phrase in self._UNCERTAINTY_PHRASES)

    def _parse_response(self, response: str) -> dict:
        default = {
            "complete": False,
            "reason": "Unable to verify completion.",
            "missing_steps": "",
        }
        if not response:
            return default

        data = self._extract_first_json_dict(response)
        if not data:
            # Be conservative: non-JSON output is treated as incomplete.
            return {
                "complete": False,
                "reason": "Verifier response was not valid JSON; treating as incomplete.",
                "missing_steps": response.strip()[:500],
            }

        complete = self._to_bool(data.get("complete", data.get("done", False)))
        reason = str(data.get("reason", data.get("analysis", ""))).strip()
        missing_steps = self._normalize_missing_steps(
            str(data.get("missing_steps", data.get("next", "")))
        )

        # If the verifier expresses uncertainty, do not accept completion.
        # If it claims completion but also lists missing steps, do not accept completion.
        if complete and (missing_steps or self._contains_uncertainty(reason)):
            complete = False
            if not reason:
                reason = "Marked complete but did not provide clear evidence."
            if not missing_steps:
                missing_steps = "Provide clear on-screen confirmation that all task requirements are satisfied."

        return {
            "complete": complete,
            "reason": reason or default["reason"],
            "missing_steps": "" if complete else missing_steps,
        }

    def verify_completion(
        self,
        instruction: str,
        observation: dict,
        trajectory: str = "",
        executor_plan: str | None = None,
    ) -> dict:
        truncated_traj = trajectory if len(trajectory) <= 3000 else trajectory[-3000:]

        obs_text_parts = []
        if isinstance(observation, dict):
            # Include lightweight, textual state when available (helps reduce false positives).
            for key in (
                "accessibility_tree",
                "a11y_tree",
                "ocr",
                "active_app",
                "window_title",
            ):
                value = observation.get(key)
                if value:
                    value_str = str(value)
                    value_str = (
                        value_str
                        if len(value_str) <= 1500
                        else value_str[:1500] + "..."
                    )
                    obs_text_parts.append(f"{key}: {value_str}")
        obs_text = "\n".join(obs_text_parts)

        # Self-categorize the task and build verification hints.
        # Check for infeasibility signals in BOTH the trajectory AND the
        # executor plan (which contains the agent's reasoning, not just action code).
        combined_text_for_signals = truncated_traj
        if executor_plan:
            combined_text_for_signals += "\n" + str(executor_plan)
        categories = self._categorize_task(instruction, truncated_traj)
        infeasibility = self._check_infeasibility_signals(combined_text_for_signals)
        category_hints = self._build_category_hints(categories, infeasibility)
        hint_block = ""
        if category_hints:
            hint_block = (
                "\nVerification guidance for this task type "
                "(IMPORTANT - how the task will actually be graded by the user):\n"
                f"{category_hints}\n"
            )

        prompt = textwrap.dedent(
            f"""\
            Determine if the task is fully completed RIGHT NOW.
            Be strict: if you cannot directly verify completion from the screenshot and provided context, mark incomplete.
            Task: {instruction}
            {hint_block}Latest plan or action: {executor_plan or "N/A"}
            Recent trajectory (truncated): {truncated_traj or "N/A"}
            Other observation text (truncated): {obs_text or "N/A"}
            Respond in the requested JSON format.
            """
        )

        self.agent.reset()
        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="verifier")
        verdict = self._parse_response(response)
        logger.info("Verifier verdict: %s | raw: %s", verdict, response)
        return verdict
