from datetime import datetime
from typing import List, Dict, Optional, Any, Union
from guardianhub import get_logger
from guardianhub.agents.services.episodic_manager import EpisodicManager
from guardianhub.clients import VectorClient, GraphDBClient, ToolRegistryClient
from guardianhub.clients.llm_client import LLMClient
# 🎯 IMPORT OUR NEW CONTRACTS
from guardianhub.models import VectorQueryRequest, VectorQueryResponse
from guardianhub.models.common.common import KeywordList

logger = get_logger(__name__)


class MemoryManager:
    def __init__(
            self,
            episodic_manager: Optional[EpisodicManager] = None,

    ):
        self.llm_client = LLMClient()
        self.librarian = episodic_manager or EpisodicManager()

    # In MemoryManager.py

    async def get_reasoning_context(
            self,
            query: str,
            template_id: str,
            capabilities: Dict[str, Any]  # 🎯 Now correctly identified as AgentName: Spec dict
    ) -> Dict[str, Any]:

        # 1. Extraction for Filtering (The "Names")
        # Since capabilities is Dict[str, Spec], we iterate through the specs in .values()
        active_capability_names = {
            spec.get('name')
            for spec in capabilities.values()
            if isinstance(spec, dict) and spec.get('name')
        }

        # 2. DELEGATED RETRIEVAL
        keywords_list = await self._resolve_keywords(query, list(capabilities.values()))
        facts = await self.librarian.get_infrastructure_facts(template_id, keywords_list)
        episodes = await self.librarian.get_recent_episodes(query, template_id)

        # 3. BELIEF FILTERING (Dharma)
        raw_beliefs = await self.librarian.get_ace_lessons(query=query, keywords=keywords_list)
        # We filter lessons by checking if the 'related_tool' matches a capability 'name'
        beliefs = self._filter_actionable_beliefs(raw_beliefs, active_capability_names)

        return {
            "facts": facts,
            "beliefs": beliefs,
            "episodes": episodes,
            "agentic_capabilities": capabilities,  # 🚀 The full Spec-Map for LLM reasoning
            "metadata": {
                "assembled_at": datetime.utcnow().isoformat(),
                "active_missions": list(active_capability_names)
            }
        }

    def _filter_actionable_beliefs(self, beliefs: List[Dict], active_capabilities: set) -> List[Dict]:
        """
        Filters ACE Lessons to ensure the agent only 'remembers' strategies
        aligned with its current Agentic Capabilities.
        """
        actionable = []
        for b in beliefs:
            # A belief is actionable if it's a global strategy OR linked to a tool we currently possess
            is_actionable = b.get("related_tool") in active_capabilities or b.get("type") == "STRATEGY"
            b["is_actionable"] = is_actionable
            actionable.append(b)
        return actionable

    # --- 🛰️ INTERNAL TIERED FETCHERS ---

    async def _resolve_keywords(self, query: str, tools: List[Dict]) -> List[str]:
        """Resolves keywords from tool metadata or LLM extraction."""
        # 1. Try Tool Metadata first (High Signal)
        if tools and isinstance(tools[0], dict):
            excerpt = tools[0].get('metadata', {}).get('llama_index_metadata', {}).get('excerpt_keywords')
            if excerpt:
                kw = [k.strip() for k in excerpt.split(',')] if isinstance(excerpt, str) else excerpt
                return kw

        # 2. Fallback to LLM extraction (Cognitive Skimming)
        llm_result = await self._extract_search_keywords(query)
        keywords = llm_result.keywords if llm_result.keywords else []
        # 2. 🎯 TACTICAL ADDITION: Inject Resource Clues
        # If the query mentions part of a known resource, add the whole resource name
        # For now, let's just make sure "database" triggers "db" matches.
        if "database" in query.lower():
            keywords.append("db")
        return list(set(keywords))  # Deduplicate

    async def _extract_search_keywords(self, query: str) -> KeywordList:
        """LLM-powered technical entity extraction."""
        # 🎯 ALIGNMENT: Match the prompt narrative to the Pydantic Model structure
        system_prompt = """
            You are a technical entity extractor. 
            Analyze the infrastructure request and extract searchable technical entities.
            """
        try:
            return await self.llm_client.invoke_structured_model(
                user_input=query,
                system_prompt_template=system_prompt,
                response_model_name="KeywordList",
            )
        except Exception:
            return KeywordList(keywords=self._basic_keyword_extraction(query), source="fallback")

    def _basic_keyword_extraction(self, query: str) -> List[str]:
        """Regex-free clean split fallback."""
        noise = {"need", "check", "show", "status", "what", "find"}
        words = query.lower().replace("?", "").split()
        return [w for w in words if w not in noise and len(w) > 2][:5]

    async def assimilate_proposal(self, current_metadata: Dict, proposal_payload: Dict) -> Dict:
        """
        [DHARMA INGESTION]
        Standardizes how the Architect's plan enters the nervous system.
        Ensures the 'Why' (Reflection) is paired with the 'What' (MacroPlan).
        """
        session_id = proposal_payload.get('session_id', 'unknown')
        logger.info(f"🧠 [MEMORY] Assimilating Tactical Proposal for session {session_id}")

        raw_steps = proposal_payload.get("macro_plan", [])
        sanitized_plan = []

        # 🎯 FIX: Force 'pending' status and ensure technical IDs exist
        for i, step in enumerate(raw_steps):
            s_copy = step.copy()
            s_copy["status"] = "pending"  # 🛡️ SHIELD: Never let a plan start as 'completed'
            if not s_copy.get("step_id"):
                s_copy["step_id"] = f"step_{i + 1}_{datetime.utcnow().strftime('%M%S')}"
            sanitized_plan.append(s_copy)

        # 1. Identity & DNA Mapping
        # We preserve the original DNA to ensure the 'Law' remains consistent
        dna_anchor = current_metadata.get("mission_dna") or "Audit Only"
        template_id = current_metadata.get("template_id", "TPL-GENERIC")

        # 2. Extract Narrative Logic
        # This is the Specialist's ephemeral reasoning that will eventually become a 'Lesson'
        reflection = proposal_payload.get("reflection") or proposal_payload.get("rationale",
                                                                                "No tactical rationale provided.")


        return {
            "macro_plan": sanitized_plan,
            "metadata": {
                **current_metadata,
                "handshake_status": "PLAN_READY",
                "dispatch_in_progress": False,
                "law_segment": dna_anchor,  # Permanent Governance
                "wisdom_segment": reflection,  # Decision Logic for future Hindsight
                "template_id": template_id,
                "ingested_at": datetime.utcnow().isoformat()
            }
        }

    async def assimilate_intervention(
            self,
            current_plan: List,
            current_context: Dict,
            current_metadata: Dict,
            result_payload: Dict
    ) -> Dict:
        """
        [SATYA CONSOLIDATION]
        Standardizes how real-world tool results are merged into the Mission State.
        Prepares the data for the 'After Action Report' (AAR).
        """
        mission_id = result_payload.get('mission_id', 'unknown')
        logger.info(f"⚡ [MEMORY] Assimilating intervention results for mission {mission_id}")

        # 1. Plan State Update
        # We find the 'running' step and mark it complete with real data
        updated_plan = []
        for step in current_plan:
            new_step = step.copy()
            if new_step.get("status") == "running" or new_step.get("step_id") == result_payload.get("step_id"):
                new_step["status"] = "completed"
                # We store the raw output here for technical tracing
                new_step["result_data"] = result_payload.get("result")
                new_step["completed_at"] = datetime.utcnow().isoformat()
            updated_plan.append(new_step)

        # 2. Context Enrichment
        # We add the result to the macro_context so the next step's LLM prompt has it
        new_context = {**current_context}
        new_context[mission_id] = result_payload

        # 3. Final State Delta
        return {
            "macro_plan": updated_plan,
            "macro_context": new_context,
            "metadata": {
                **current_metadata,
                "dispatch_in_progress": False,  # Release the workflow lock
                "last_mission_id": mission_id,
                "last_result_status": result_payload.get("status", "unknown")
            }
        }