import json
import logging
from typing import Any, Dict, Optional

from GramAddict.core.perception.spatial_parser import SpatialNode

logger = logging.getLogger(__name__)


def _parse_yes_no(response: str) -> Optional[bool]:
    """Parses a VLM response to find a definitive YES or NO without substring-matching 'not' or 'now'."""
    text = response.strip()

    # Try parsing as JSON first
    if text.startswith("{"):
        try:
            data = json.loads(text)
            for k, v in data.items():
                if str(k).strip().upper() == "YES" or str(v).strip().upper() == "YES":
                    return True
                if str(k).strip().upper() == "NO" or str(v).strip().upper() == "NO":
                    return False
                if str(k).strip().lower() == "success" and isinstance(v, bool):
                    return v

            # If it is valid JSON but we couldn't definitively find YES/NO,
            # do NOT fall through to text matching
            return None
        except Exception:
            # Prevent JSON parsing fall-throughs
            return None

    text_lower = text.lower()
    if text_lower.startswith("yes"):
        return True
    if text_lower.startswith("no") and not text_lower.startswith("now") and not text_lower.startswith("not"):
        return False

    return None


# ═══════════════════════════════════════════════════════
# Semantic Match Keywords — SSOT for intent → element validation
# ═══════════════════════════════════════════════════════

# Maps toggle-intent keywords to required element markers.
# If the intent contains the key, the clicked element MUST
# contain at least one of the corresponding markers in its
# text, content_desc, or resource_id.
# ZERO MAINTENANCE: Only English words and resource_id fragments allowed.
# No localized strings — the bot must work on any device language.
TOGGLE_INTENT_MARKERS = {
    "follow": ["follow", "button_follow"],
    "like": ["like", "heart", "button_like"],
    "save": ["save", "saved", "bookmark"],
}


class ActionMemory:
    """
    Handles the caching, tracking, and negative reinforcement (unlearning) of UI interactions.
    Decouples the memory layer from the core parsing engine.
    """

    def __init__(self, ui_memory=None):
        # We optionally inject UIMemoryDB to decouple tests
        if ui_memory is None:
            from GramAddict.core.qdrant_memory import UIMemoryDB

            self.ui_memory = UIMemoryDB()
        else:
            self.ui_memory = ui_memory

        self._last_click_context: Optional[Dict[str, Any]] = None

    def track_click(self, intent: str, node: SpatialNode, xml_context: str = ""):
        """Stores the context of a click before it's actually performed."""
        semantic_string = f"text: '{node.text}', desc: '{node.content_desc}', id: '{node.resource_id}'"

        self._last_click_context = {
            "intent": intent,
            "node_dict": node.to_dict(),
            "semantic_string": semantic_string,
            "xml_context": xml_context,
        }
        logger.debug(f"🧠 [ActionMemory] Tracking tentative click for intent: '{intent}' -> {semantic_string}")

    def confirm_click(self, intent: str = None):
        """Positive Reinforcement: Confirms the last click was successful.

        Guard: Refuses to store in Qdrant if the clicked element does not
        semantically match the intent. Prevents memory poisoning.
        """
        ctx = self._last_click_context
        if not ctx:
            return

        if intent and ctx["intent"] != intent:
            return

        # ── Semantic Mismatch Guard ──
        if not _intent_matches_node(ctx["intent"], ctx["semantic_string"]):
            logger.warning(
                f"🛡️ [ActionMemory] BLOCKED confirm_click for '{ctx['intent']}' — "
                f"clicked element does not match intent: {ctx['semantic_string']}"
            )
            self._last_click_context = None
            return

        logger.info(
            f"✅ [ActionMemory] Confirming success for '{ctx['intent']}'. Boosting confidence.",
            extra={"color": "\x1b[32m"},
        )

        # Store or boost in Qdrant
        try:
            # Check if it exists first
            existing = self.ui_memory.retrieve_memory(ctx["intent"], ctx["xml_context"])
            if existing:
                self.ui_memory.boost_confidence(ctx["intent"], ctx["xml_context"])
            else:
                self.ui_memory.store_memory(ctx["intent"], ctx["xml_context"], ctx["node_dict"])
        except Exception as e:
            logger.warning(f"Failed to confirm click in Qdrant: {e}")

        self._last_click_context = None

    def reject_click(self, intent: str = None):
        """Negative Reinforcement: Penalizes a failed click (Unlearning)."""
        ctx = self._last_click_context
        if not ctx:
            return

        if intent and ctx["intent"] != intent:
            return

        logger.warning(
            f"❌ [ActionMemory] Click failed for '{ctx['intent']}'. Applying penalty.", extra={"color": "\x1b[31m"}
        )

        try:
            self.ui_memory.decay_confidence(ctx["intent"], ctx["xml_context"])
        except Exception as e:
            logger.warning(f"Failed to decay confidence in Qdrant: {e}")

        self._last_click_context = None

    def verify_success(
        self, intent: str, pre_click_xml: str, post_click_xml: str, device=None, confidence: float = 0.0
    ) -> Optional[bool]:
        """
        Structural and Visual verification: Did the UI actually change after the click?
        """
        intent_lower = intent.lower()
        post_xml_lower = post_click_xml.lower()

        # Specific check for opening a post (from explore/profile grid)
        if "view a post" in intent_lower or "first image" in intent_lower or "grid item" in intent_lower:
            if (
                "row_feed_photo_imageview" in post_xml_lower
                or "row_feed_button_like" in post_xml_lower
                or "clips_viewer_view_pager" in post_xml_lower
            ):
                return True
            if (
                "explore_action_bar" in post_xml_lower
                and "row_feed_button_like" not in post_xml_lower
                and "clips_viewer" not in post_xml_lower
            ):
                logger.warning(f"⚠️ [ActionMemory] Still on grid after trying to '{intent}'. Verification FAIL.")
                return False  # Still on grid, definitely failed

        # Specific check for opening a profile
        if "profile" in intent_lower or "author" in intent_lower or "username" in intent_lower:
            if "profile_header_container" in post_xml_lower:
                logger.info("✅ [ActionMemory] Structural check confirmed profile navigation success.")
                return True
            else:
                logger.warning(
                    f"⚠️ [ActionMemory] Profile header NOT found after trying to '{intent}'. Verification FAIL."
                )
                return False

        # Specific check for navigating to Home Feed
        if "home feed" in intent_lower or "home tab" in intent_lower:
            if "main_feed_action_bar" in post_xml_lower:
                logger.info("✅ [ActionMemory] Structural check confirmed Home Feed navigation success.")
                return True

        # Specific check for navigating to Explore Feed
        if "explore feed" in intent_lower or "explore tab" in intent_lower or "search" in intent_lower:
            if "explore_action_bar" in post_xml_lower or "action_bar_search_edit_text" in post_xml_lower:
                logger.info("✅ [ActionMemory] Structural check confirmed Explore Feed navigation success.")
                return True

        state_toggles = ["like", "save", "follow", "heart"]
        is_toggle = any(t in intent_lower for t in state_toggles)

        # ── VLM Verification Fallback ──

        # If we are highly confident (e.g. pulled from Qdrant memory), bypass heavy VLM
        if device and confidence < 0.95:
            logger.info(
                f"👁️ [ActionMemory] Confidence ({confidence:.2f}) < 0.95. Handing over verification for '{intent}' to VLM visual analysis..."
            )
            from GramAddict.core.perception.semantic_evaluator import SemanticEvaluator

            evaluator = SemanticEvaluator()

            # Build context of what was actually clicked
            clicked_context = ""
            if self._last_click_context:
                clicked_context = f"The element that was tapped: {self._last_click_context['semantic_string']}. "

            # Ask VLM to be the absolute source of truth
            prompt = (
                f"The user just attempted to perform the action: '{intent}'. "
                f"{clicked_context}"
                f"Look at the current screen carefully. Was the action successful? "
            )
            if is_toggle:
                prompt += (
                    "If the intent was 'follow', does the button now indicate 'Following' or 'Requested'? "
                    "If it was 'like', is the heart icon clearly active/red? "
                    "If the screen shifted completely to a profile when you just wanted to like/follow from a feed, it FAILED. "
                    "If the tapped element does NOT sound like a like/follow button (e.g. it's a caption, comment field, or post content), it FAILED. "
                )
            else:
                prompt += (
                    f"Does the current screen match the expected outcome of '{intent}'? "
                    f"For example, if the intent was to open a post/photo, are you looking at a post view (not a user profile or story)? "
                    f"If the intent was to open a profile, are you on a profile page? "
                    f"If the intent was to go back, are you on the previous screen? "
                )
            prompt += "Answer ONLY with the word YES or NO."

            try:
                screenshot = device.get_screenshot_b64()
                if not screenshot:
                    raise ValueError("No screenshot available from device")
                response = evaluator._query_vlm(prompt, screenshot)

                decision = _parse_yes_no(response) if response else None

                if decision is True:
                    logger.debug(f"🧠 [ActionMemory] VLM visually confirmed success for '{intent}'.")
                    return True
                elif decision is False:
                    logger.warning(
                        f"⚠️ [ActionMemory] VLM visual verification FAILED for '{intent}'. VLM replied: '{response}'"
                    )
                    return False
                else:
                    # VLM returned ambiguous response (JSON, mixed signals, etc.)
                    # Don't treat as hard failure — fall through to structural delta verification
                    logger.info(
                        f"🧠 [ActionMemory] VLM response for '{intent}' was not YES/NO "
                        f"(got: '{response[:80]}...'). Falling through to structural verification."
                    )
            except Exception as e:
                logger.error(f"Failed to query VLM for visual verification: {e}")
                # Fallthrough to structural delta if VLM crashes

        # ── Pre-Structural Semantic Gate ──
        if is_toggle and self._last_click_context:
            if not _intent_matches_node(intent, self._last_click_context["semantic_string"]):
                logger.warning(
                    f"🛡️ [ActionMemory] Semantic mismatch: '{intent}' does not match "
                    f"clicked element {self._last_click_context['semantic_string']}. Verification FAIL."
                )
                return False

        # Fallback to structural delta
        logger.info(f"DEBUG: len(pre_click_xml)={len(pre_click_xml)} len(post_click_xml)={len(post_click_xml)}")
        diff = abs(len(pre_click_xml) - len(post_click_xml))
        logger.info(f"DEBUG: diff={diff}")

        if is_toggle:
            if diff > 1000:
                logger.warning(
                    f"⚠️ [ActionMemory] Massive structural shift ({diff} chars) for state-toggle '{intent}'. Navigated away by mistake? Verification FAIL."
                )
                return False
            if diff > 0:
                logger.debug(f"🧠 [ActionMemory] Structural delta detected for toggle '{intent}'. Verification PASS.")
                return True
            logger.warning(
                f"⚠️ [ActionMemory] Zero structural shift (diff={diff}) for state-toggle '{intent}'. Verification FAIL."
            )
            return False
            # If the intent is an abstract goal (like "find customers"), diff > 50 is NOT enough.
            # We must force visual VLM confirmation because clicking the wrong thing (like "Create highlight")
            # also produces a large diff but achieves the wrong goal.
            if diff > 50:
                # Is it a standard structural transition?
                from GramAddict.core.screen_topology import ScreenTopology

                # We don't have screen type here, so we just check if it's in the HD Map keys
                logger.info(f"DEBUG: intent is '{intent}'")
                logger.info(
                    f"DEBUG: TRANSITIONS keys are: {[list(t.keys()) for t in ScreenTopology.TRANSITIONS.values()]}"
                )
                is_standard = any(intent in transitions for transitions in ScreenTopology.TRANSITIONS.values())
                logger.info(f"DEBUG: is_standard={is_standard}")

                if is_standard:
                    logger.debug(
                        f"🧠 [ActionMemory] Structural change detected for known navigation '{intent}'. Verification PASS."
                    )
                    return True
                else:
                    logger.info(
                        f"👁️ [ActionMemory] Abstract intent '{intent}' caused UI change. Forcing VLM visual verification..."
                    )
                    # For abstract intents, we must visually verify if it actually helped!
                    # If device is available, we use VLM. If not, we fail safe.
                    if device:
                        from GramAddict.core.perception.semantic_evaluator import SemanticEvaluator

                        evaluator = SemanticEvaluator()
                        prompt = f"The user just attempted to perform the action: '{intent}'. Does the current screen match the expected outcome? Answer ONLY with the word YES or NO."
                        try:
                            response = evaluator._query_vlm(prompt, device.get_screenshot_b64())
                            decision = _parse_yes_no(response) if response else None
                            if decision is True:
                                return True
                            else:
                                logger.warning(
                                    f"⚠️ [ActionMemory] VLM rejected success for abstract intent '{intent}'. Response: '{response}'"
                                )
                                return False
                        except Exception as e:
                            logger.error(f"VLM visual verification failed: {e}")

                    logger.warning(f"⚠️ [ActionMemory] Cannot visually verify abstract intent '{intent}'. Failing safe.")
                    return False

            # If diff <= 50 for non-toggle
            logger.warning(
                f"⚠️ [ActionMemory] Insufficient structural change (diff={diff}) for non-toggle '{intent}'. Verification FAIL."
            )
            return False


def _intent_matches_node(intent: str, semantic_string: str) -> bool:
    """Checks if the clicked element semantically matches the toggle intent.

    For toggle intents (follow, like, save), the clicked element MUST contain
    at least one of the required keywords in its text/desc/id. This prevents
    photo grid items, captions, and other unrelated elements from being
    falsely confirmed as successful interactions.

    For non-toggle intents, returns True (no restriction).
    """
    intent_lower = intent.lower()
    semantic_lower = semantic_string.lower()

    for intent_keyword, required_markers in TOGGLE_INTENT_MARKERS.items():
        if intent_keyword in intent_lower:
            if any(marker in semantic_lower for marker in required_markers):
                return True
            logger.debug(
                f"🛡️ [SemanticGuard] Intent '{intent}' requires markers "
                f"{required_markers} but element has: {semantic_string}"
            )
            return False

    # Non-toggle intents pass through
    return True
