import logging
import time
from typing import List, Optional

from GramAddict.core.perception.screen_identity import ScreenType
from GramAddict.core.qdrant_memory import QdrantBase

logger = logging.getLogger(__name__)


class NavigationKnowledge:
    """
    Manages the bot's learned understanding of the Instagram UI.
    Discovered dynamically through exploration and success.
    """

    def __init__(self, username: str):
        self.username = username
        try:
            self._db = QdrantBase("navigation_knowledge", vector_size=768)
        except Exception:
            self._db = None

        # In-memory cache for rapidly avoiding traps during exploration
        # In-memory cache for rapidly avoiding traps during exploration
        self._learned_screen_mappings = {}
        self._learned_traps = set()

    def wipe(self):
        """Wipe all learned knowledge from Qdrant."""
        if self._db and self._db.is_connected:
            try:
                self._db.wipe_collection()
            except Exception as e:
                logger.warning(f"⚠️ [NavigationKnowledge] Could not wipe knowledge: {e}")

    def update_username(self, username: str):
        """Update username and reconnect DB if needed."""
        if self.username != username:
            self.username = username
            try:
                self._db = QdrantBase("navigation_knowledge", vector_size=768)
            except Exception:
                self._db = None

    def get_requirements(self, goal: str) -> List[ScreenType]:
        """Get required screens for a goal. Returns known requirements or empty list."""
        if not self._db or not self._db.is_connected:
            return []

        try:
            from qdrant_client.models import FieldCondition, Filter, MatchValue

            results = self._db.client.scroll(
                collection_name=self._db.collection_name,
                scroll_filter=Filter(must=[FieldCondition(key="goal", match=MatchValue(value=goal))]),
                limit=1,
            )[0]
            if results:
                screen_name = results[0].payload.get("required_screen")
                logger.debug(f"🧠 [Nav Knowledge] Found requirement for '{goal}': {screen_name}")
                if screen_name:
                    return [ScreenType[screen_name]]
        except Exception as e:
            logger.warning(f"⚠️ [Nav Knowledge] Search error: {e}")
        return []

    def learn_goal_requirement(self, goal: str, screen_type: ScreenType):
        """Learn that achieving 'goal' lands us on 'screen_type'."""
        if not self._db or not self._db.is_connected:
            logger.warning("⚠️ [Nav Knowledge] Cannot learn: DB not connected")
            return

        seed = f"req_{goal}"
        vec = self._db._get_embedding(f"goal_requirement: {goal}")
        payload = {"goal": goal, "required_screen": screen_type.name, "timestamp": time.time()}
        self._db.upsert_point(seed, payload, vector=vec)
        logger.info(f"🧠 [Nav Knowledge] Learned: '{goal}' → {screen_type.name}")

    def get_action_for_screen(self, target_screen: ScreenType) -> Optional[str]:
        """Find which action leads to this screen."""
        for action, screen in self._learned_screen_mappings.items():
            if screen == target_screen:
                return action

        if not self._db or not self._db.is_connected:
            return None

        try:
            from qdrant_client.models import FieldCondition, Filter, MatchValue

            results = self._db.client.scroll(
                collection_name=self._db.collection_name,
                scroll_filter=Filter(
                    must=[FieldCondition(key="result_screen", match=MatchValue(value=target_screen.name))]
                ),
                limit=1,
            )[0]
            if results:
                return results[0].payload.get("action")
        except Exception:
            pass
        return None

    def get_screen_for_action(self, action: str) -> Optional[ScreenType]:
        """Find where this action leads to to avoid looping traps."""
        if action in self._learned_screen_mappings:
            return self._learned_screen_mappings[action]

        if not self._db or not self._db.is_connected:
            return None

        try:
            from qdrant_client.models import FieldCondition, Filter, MatchValue

            results = self._db.client.scroll(
                collection_name=self._db.collection_name,
                scroll_filter=Filter(must=[FieldCondition(key="action", match=MatchValue(value=action))]),
                limit=1,
            )[0]
            if results:
                screen_name = results[0].payload.get("result_screen")
                if screen_name:
                    return ScreenType[screen_name]
        except Exception:
            pass
        return None

    def learn_screen_mapping(self, action: str, result_screen: ScreenType):
        """Learn that taking 'action' leads to 'result_screen'."""
        if not self._db or not self._db.is_connected:
            return

        seed = f"map_{action}"
        vec = self._db._get_embedding(f"screen_mapping: {result_screen.name}")
        payload = {"action": action, "result_screen": result_screen.name, "timestamp": time.time()}

        self._learned_screen_mappings[action] = result_screen

        self._db.upsert_point(seed, payload, vector=vec)
        logger.info(f"🧠 [Nav Knowledge] Learned Mapping: '{action}' → {result_screen.name}")

    def get_screen_for_tab(self, tab_id: str) -> Optional[ScreenType]:
        """Find where this tab leads to to avoid looping traps."""
        if tab_id in self._learned_screen_mappings:
            return self._learned_screen_mappings[tab_id]

        if not self._db or not self._db.is_connected:
            return None

        try:
            from qdrant_client.models import FieldCondition, Filter, MatchValue

            results = self._db.client.scroll(
                collection_name=self._db.collection_name,
                scroll_filter=Filter(must=[FieldCondition(key="tab_id", match=MatchValue(value=tab_id))]),
                limit=1,
            )[0]
            if results:
                s_name = results[0].payload.get("result_screen")
                if s_name:
                    return ScreenType[s_name]
        except Exception:
            pass
        return None

    def learn_trap(self, screen_type: ScreenType, action: str, trap_reason: str = "softlock"):
        """Aversively learn that an action on a screen is dangerous/useless."""
        trap_key = f"{screen_type.name}_{action}"
        self._learned_traps.add(trap_key)

        if not self._db or not self._db.is_connected:
            return

        seed = f"trap_{trap_key}"
        # Aversive vector is completely orthogonal to normal goals to prevent retrieval overlap
        vec = self._db._get_embedding(f"trap_avoidance: {trap_key} {trap_reason}")
        payload = {
            "trap_screen": screen_type.name,
            "trap_action": action,
            "trap_reason": trap_reason,
            "timestamp": time.time(),
        }
        self._db.upsert_point(seed, payload, vector=vec)
        logger.error(f"💀 [Aversive Learning] BURNED action '{action}' on {screen_type.name} due to: {trap_reason}")

    def is_trap(self, screen_type: ScreenType, action: str) -> bool:
        """Check if an action on this screen is a known trap."""
        trap_key = f"{screen_type.name}_{action}"
        if trap_key in self._learned_traps:
            return True

        if not self._db or not self._db.is_connected:
            return False

        try:
            from qdrant_client.models import FieldCondition, Filter, MatchValue

            results = self._db.client.scroll(
                collection_name=self._db.collection_name,
                scroll_filter=Filter(
                    must=[
                        FieldCondition(key="trap_screen", match=MatchValue(value=screen_type.name)),
                        FieldCondition(key="trap_action", match=MatchValue(value=action)),
                    ]
                ),
                limit=1,
            )[0]
            if results:
                self._learned_traps.add(trap_key)
                return True
        except Exception:
            pass
        return False
