from typing import List, Dict, Optional, Any
from guardianhub import get_logger
from guardianhub.clients.llm_client import LLMClient
from guardianhub.config.settings import settings
from guardianhub.models.common.common import KeywordList, KeyValuePair, ErrorResponse, Status

logger = get_logger(__name__)

class MemoryManager:
    def __init__(
        self, 
        vector_client, 
        graph_client, 
        tool_registry, 
    ):
        self.vector_client = vector_client
        self.graph_client = graph_client
        self.tool_registry = tool_registry
        self.llm_client = LLMClient()

    async def get_reasoning_context(self, query: str, template_id: str, tools):
        """
        Retrieve reasoning context including facts, beliefs, and episodes.

        Args:
            query: The search query to filter context.

        Returns:
            Dict containing facts, beliefs, and episodes.
        """
        logger.info(f"Starting context retrieval for query: {query}")

        # 1. Fetch available tools
        logger.info("Fetching active tools...")
        try:
            # active_tools = await self.tool_registry.search_dynamic_tool(query)
            current_caps = {t['name'] for t in tools}
            logger.info(f"Found {len(tools)} active tools: {current_caps}")
            if tools and isinstance(tools[0], dict) and tools[0].get('metadata', {}).get('llama_index_metadata', {}).get('excerpt_keywords'):
                # Use excerpt_keywords from the first tool's llama_index_metadata
                keywords_list = tools[0].get('metadata', {}).get('llama_index_metadata', {}).get('excerpt_keywords')
                if isinstance(keywords_list, str):
                    # If it's a string, split by commas and clean up
                    keywords_list = [k.strip() for k in keywords_list.split(',') if k.strip()]
                logger.info(f"Using {len(keywords_list)} keywords from tool metadata: {keywords_list}")
            else:
                # Fallback to the original keyword extraction if no excerpt_keywords available
                keywords = await self._extract_search_keywords(query)
                keywords_list = keywords.keywords if hasattr(keywords, 'keywords') and keywords.keywords else []
                logger.info(f"Extracted {len(keywords_list)} keywords: {keywords_list}")

        except Exception as e:
            logger.error(f"Error fetching active tools: {str(e)}")
            active_tools = []
            current_caps = set()
            keywords_list = []

        # 2. Fetch Tiered Memory
        facts, all_beliefs, episodes = [], [], []

        # Fetch facts
        try:
            logger.info("Fetching infrastructure facts...")
            facts = await self.graph_client.get_infrastructure_facts(template_id=template_id,keywords=keywords_list)
            logger.info(f"Retrieved {len(facts)} facts")
        except Exception as e:
            logger.error(f"Error fetching facts: {str(e)}")
            facts = []

        # Fetch beliefs
        try:
            logger.info("Fetching ACE lessons...")

            all_beliefs = await self.graph_client.get_ace_lessons(query=query,keywords=keywords_list)
            logger.info(f"Retrieved {len(all_beliefs)} beliefs")
        except Exception as e:
            logger.error(f"Error fetching beliefs: {str(e)}")
            all_beliefs = []

        # Fetch episodes
        try:
            logger.info("Fetching recent episodes...")
            episodes = await self.vector_client.get_recent_episodes(query)
            logger.info(f"Retrieved {len(episodes)} episodes")
        except Exception as e:
            logger.error(f"Error fetching episodes: {str(e)}")
            episodes = []

        # 3. Categorize and Filter "Free Flowing" Beliefs
        logger.info("Categorizing beliefs...")
        actionable_beliefs = []
        for b in all_beliefs:
            try:
                related_tool = b.get("related_tool", "")
                belief_type = b.get("type", "")
                is_actionable = related_tool in current_caps or belief_type == "STRATEGY"

                if isinstance(b, dict):
                    b["actionable"] = is_actionable
                    actionable_beliefs.append(b)
                else:
                    logger.warning(f"Unexpected belief format: {b}")
            except Exception as e:
                logger.error(f"Error processing belief {b}: {str(e)}")

        logger.info(
            f"Context retrieval complete. Found {len(facts)} facts, {len(actionable_beliefs)} actionable beliefs, and {len(episodes)} episodes.")

        return {
            "facts": facts,
            "beliefs": actionable_beliefs,
            "episodes": episodes
        }

    async def _extract_search_keywords(self, query: str) -> KeywordList:
        """
        Extract relevant search keywords using LLM-powered extraction with fallback.
        
        Args:
            query: The input query to extract keywords from
            
        Returns:
            KeywordList containing extracted keywords and metadata
        """
            # Try LLM-based extraction first
        system_prompt = """
                You are a technical entity extractor for an IT Knowledge Graph.
                Extract key infrastructure terms, tool names, or CI types from the input.
                Return ONLY a JSON list of strings.
                Example: "I need to check my cmdb health" -> ["cmdb", "health"]
                """
        try:
            result = await self.llm_client.invoke_structured_model(
                model_key="default",
                user_input=query,
                system_prompt_template=system_prompt,
                response_model_name="KeywordList",
            )

            if result and result.keywords:
                logger.info(
                    f"Extracted {len(result.keywords)} keywords using LLM",
                    extra={
                        "keywords": result.keywords,
                        "confidence": result.confidence
                    }
                )
                return result

        except Exception as llm_error:
            logger.warning(
                "LLM-based keyword extraction failed, falling back to basic extraction",
                extra={"error": str(llm_error)}
            )

            return KeywordList(
                keywords=[""],
                source="basic",
                confidence=0.5
            )


    
    def _basic_keyword_extraction(self, query: str) -> List[str]:
        """Basic keyword extraction as fallback when LLM extraction fails."""
        noise = {
            "need", "check", "please", "show", "find", "some", "any", "status", 
            "the", "what", "how", "when", "where", "why", "which"
        }
        words = query.lower().replace("?", "").replace("!", "").split()
        return [w for w in words if w not in noise and len(w) > 2][:5]

    async def _extract_search_keywords_later(self, query: str) -> List[str]:
        """
        Cognitive Skimming: Uses the LLM to extract technical entities
        and infrastructure keywords from conversational user input.
        """
        system_prompt = """
        You are a technical entity extractor for an IT Knowledge Graph.
        Extract key infrastructure terms, tool names, or CI types from the input.
        Return ONLY a JSON list of strings.
        Example: "I need to check my cmdb health" -> ["cmdb", "health"]
        """

        try:
            # Using your existing LLM service to get a structured list
            keywords = await self.llm_client.invoke_structured_model(
                system_prompt_template=system_prompt,
                user_input=query,
                response_model_name="KeywordList"  # Define a simple Pydantic model with List[str]
            )
            # Fallback to query if extraction is empty
            return keywords if keywords else query.lower().split()

        except Exception as e:
            logger.warning(f"Keyword extraction failed, falling back to split: {e}")
            return query.lower().replace("?", "").split()


class HindsightSimulator:
    @staticmethod
    def calculate_step_risk(step_name: str, episodes: List[Dict]) -> Dict:
        """Analyzes past episodes to predict if a step is likely to fail."""
        past_failures = [e for e in episodes if step_name in e['failed_steps']]

        if len(past_failures) > 0:
            return {
                "risk": "High",
                "suggested_buffer": 2.0,  # Suggest 2x timeout
                "reason": f"Failed in {len(past_failures)} recent episodes."
            }
        return {"risk": "Low", "suggested_buffer": 1.0}