import logging
from typing import Any, Dict, Optional

from GramAddict.core.perception.spatial_parser import SpatialNode

logger = logging.getLogger(__name__)


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."""
        ctx = self._last_click_context
        if not ctx:
            return

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

        logger.info(f"✅ [ActionMemory] Confirming success for '{ctx['intent']}'. Boosting confidence.")

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

        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?
        """
        # Specific check for explore grid
        if "first image in explore grid" in intent or "grid item" in intent:
            if "row_feed_photo_imageview" in post_click_xml or "row_feed_button_like" in post_click_xml:
                return True
            if "explore_action_bar" in post_click_xml and "row_feed_button_like" not in post_click_xml:
                return None  # Still on grid, inconclusive

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

        # 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()

            # Ask VLM to be the absolute source of truth
            prompt = (
                f"The user just attempted to perform the action: '{intent}'. "
                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. "
                )
            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()
                response = evaluator._query_vlm(prompt, screenshot)

                if response and "yes" in response.lower() and "no" not in response.lower():
                    logger.debug(f"🧠 [ActionMemory] VLM visually confirmed success for '{intent}'.")
                    return True
                else:
                    logger.warning(
                        f"⚠️ [ActionMemory] VLM visual verification FAILED for '{intent}'. VLM replied: '{response}'"
                    )
                    return False
            except Exception as e:
                logger.error(f"Failed to query VLM for visual verification: {e}")
                # Fallthrough to structural delta if VLM crashes

        # Fallback to structural delta if no device, VLM fails, or high confidence bypass
        diff = abs(len(pre_click_xml) - len(post_click_xml))

        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
        else:
            if diff > 50:
                logger.debug(
                    f"🧠 [ActionMemory] Structural change detected for navigation '{intent}'. Verification PASS."
                )
                return True

        logger.warning(f"⚠️ [ActionMemory] No structural change detected for '{intent}'. Verification FAIL.")
        return False
