from datetime import datetime
from typing import Dict, Any, List, Optional
from guardianhub import get_logger
from guardianhub.clients import VectorClient, GraphDBClient, ToolRegistryClient

logger = get_logger(__name__)


class EpisodicManager:
    def __init__(
            self,
            vector_client: Optional[VectorClient] = None,
            graph_client: Optional[GraphDBClient] = None,
            tool_registry: Optional[ToolRegistryClient] = None
    ):
        # 🛡️ Defensive anchoring to prevent the 'NoneType' crashes we saw
        self.vector_client = vector_client or VectorClient()
        self.graph_client = graph_client or GraphDBClient()
        self.tool_registry = tool_registry or ToolRegistryClient()

    async def record_resource_discovery(self, session_id: str, tool_results: Dict[str, Any]):
        """
        [SATYA SEGMENT]
        Promotes tool findings to permanent 'Shared Resource' nodes.
        Converts transient 'Observations' into persistent 'Facts'.
        """
        findings = tool_results.get("report", {}).get("detailed_metrics", [])
        if not findings:
            logger.info(f"ℹ️ No new resource discoveries for session {session_id}")
            return

        cypher_query = """
        // 1. Ensure the Episode exists (Self-healing anchor)
        MERGE (e:Episode {session_id: $session_id})

        WITH e
        UNWIND $findings AS finding

        // 2. MERGE the Resource by Name (The Entity itself)
        MERGE (r:Resource {name: finding.tool_name})
        SET r.resource_id = finding.tool_name, // Unique anchor
            r.last_observed = datetime(),
            r.current_status = finding.status,
            r.observed_in_session = $session_id,
            r.otel_trace_id = $trace_id

        // 3. Create the 'Fact' as a snapshot of this moment
        MERGE (f:Fact {fact_id: "assertion_" + finding.tool_name + "_" + toString(timestamp())})
        SET f.content = finding.factual_summary_text,
            f.type = 'infrastructure',
            f.last_updated = datetime(),
            f.observed_status = finding.status

        // 4. Bind the topology: Episode -> Fact -> Resource
        MERGE (e)-[:DISCOVERED]->(f)
        MERGE (f)-[:ASSERTED_FOR]->(r)
        """

        params = {
            "session_id": session_id,
            "findings": findings,
            "trace_id": tool_results.get("otel_trace_id", "manual_run")
        }

        # Calling the validated write method from your GraphDBClient
        await self.graph_client._execute_write_query(cypher_query, params)
        logger.info(f"✅ [SATYA] {len(findings)} resource discoveries anchored to Graph.")

    async def record_session_narrative(self, session_id: str, results: Dict[str, Any]):
        """
        [LEELA SEGMENT]
        Consolidates the mission outcome. Snapshots the 'Workforce' configuration
        to ensure the Episode is context-aware.
        """
        capabilities_snapshot = []
        try:
            # 🎯 The Sync: load_agentic_capabilities returns Dict[name, spec]
            agent_specs = await self.tool_registry.load_agentic_capabilities()

            # We store the 'Identity' of the agents active in the ecosystem during this session
            capabilities_snapshot = list(agent_specs.keys())

        except Exception as e:
            logger.warning(f"⚠️ Capability snapshot failure: {e}")

        query = """
        MERGE (e:Episode {session_id: $session_id})
        SET e.timestamp = datetime($timestamp),
            e.outcome = $outcome,
            e.capabilities = $capabilities,
            e.summary = $summary,
            e.mission_dna = $mission_dna

        // Link back to the Blueprint for 'Lawful' auditing
        WITH e
        MATCH (m:MissionBlueprint {template_id: $template_id})
        MERGE (e)-[:EXECUTED_AS]->(m)

        // Finalize the sensory 'Observations' linked to this session
        WITH e
        OPTIONAL MATCH (e)-[:MADE_OBSERVATION]->(o:Observation)
        SET o.finalized = true
        """

        params = {
            "session_id": session_id,
            "timestamp": datetime.utcnow().isoformat(),
            "outcome": results.get("status", "unknown"),
            "summary": results.get("reflection", "No synthesis provided."),
            "capabilities": capabilities_snapshot,  # 🏹 Now a clean list of agent names
            "mission_dna": results.get("mission_dna", "Standard Protocol"),
            "template_id": results.get("active_template_id", "TPL-GENERIC")
        }

        await self.graph_client._execute_write_query(query, params)
        logger.info(
            f"💾 [LEELA] Session {session_id} narrative finalized with {len(capabilities_snapshot)} agents in registry.")

    async def capture_episode(self, session_id: str, context: Dict[str, Any]):
        """
        [DHARMA SEGMENT]
        Consolidates Specialist 'Reflections' into the 'Dharma' (Belief System).
        Anchors wisdom to the Mission Blueprint to enable future Hindsight.
        """
        # 1. Extract the DNA anchor and the raw reflections
        template_id = context.get("active_template_id", "TPL-GENERIC")
        raw_lessons = context.get("reflection", [])

        # Ensure we are iterating over a list
        if isinstance(raw_lessons, dict):
            raw_lessons = [raw_lessons]

        if not raw_lessons:
            logger.info(f"ℹ️ No new wisdom to consolidate for session {session_id}")
            return

        # 2. Distill and Anchor Wisdom
        for lesson in raw_lessons:
            # We enforce a 'Muscle' check: only store lessons tied to actionable tools
            tool_id = lesson.get("tool_id")

            # 🎯 The Graph Anchor:
            # We ensure the create_lesson_node connects THIS lesson to THIS mission DNA
            try:
                # We pass the template_id so the GraphClient can MERGE the relationship
                success = await self.graph_client.create_lesson_node(
                    template_id=template_id,
                    topic=lesson.get("topic", "General Improvement"),
                    type=lesson.get("type", "OPERATIONAL"),
                    related_tool=tool_id,
                    content=lesson.get("insight", "No insight provided."),
                    session_id=session_id
                )

                if success:
                    logger.debug(f"✨ Wisdom anchored: {lesson.get('topic')} -> {template_id}")

            except Exception as e:
                logger.error(f"❌ Failed to anchor wisdom bullet: {e}")

        logger.info(f"🧠 [DHARMA] Consolidated {len(raw_lessons)} wisdom bullets for DNA {template_id}")

    async def get_hindsight(self, current_task: str, template_id: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        [DHARMA RETRIEVAL]
        Retrieves high-confidence lessons learned from past incarnations of this mission.
        Calibrates the current agent's risk based on historical 'Karma'.
        """
        # 🎯 The Shift: We now prioritize lessons linked to the specific Mission DNA (template_id)
        # but fallback to topic-based matching if we are in ad-hoc mode.

        query = """
        MATCH (m:MissionBlueprint)-[:GENERATED_LESSON]->(l:ACE_Lesson)

        // We filter by Template ID (Strict) OR by Task Topic (Flexible)
        WHERE (m.template_id = $template_id) OR 
              (toLower(l.topic) CONTAINS toLower($task)) OR 
              (toLower($task) CONTAINS toLower(l.topic))

        // Pull the Episode that birthed this lesson for outcome context
        OPTIONAL MATCH (e:Episode {session_id: l.derived_from})

        RETURN e.outcome AS past_outcome, 
               l.content AS advice,
               l.topic AS lesson_topic,
               l.created_at AS occurred_at
        ORDER BY l.created_at DESC LIMIT 3
        """

        params = {
            "task": current_task,
            "template_id": template_id
        }

        try:
            # Using the validated _execute_write_query or a generic read query wrapper
            results = await self.graph_client._execute_write_query(query, params)

            logger.info(f"🔮 [HINDSIGHT] Retrieved {len(results)} historical lessons for task: {current_task}")
            return results

        except Exception as e:
            logger.error(f"❌ [HINDSIGHT] Failed to retrieve past lessons: {e}")
            return []
