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:
    """
    The Sovereign Librarian: Standardizes how the AI Guardian remembers.
    Routes sensory data to Satya (Facts), Leela (Narratives), and Dharma (Wisdom).
    """

    def __init__(
            self,
            vector_client: Optional[VectorClient] = None,
            graph_client: Optional[GraphDBClient] = None,
            tool_registry: Optional[ToolRegistryClient] = None
    ):
        self.vector_client = vector_client or VectorClient()
        self.graph_client = graph_client or GraphDBClient()
        self.tool_registry = tool_registry or ToolRegistryClient()

    # ============================================================================
    # 1. SATYA (Truth): Sensory Infrastructure Discovery
    # ============================================================================

    async def record_resource_discovery(self, session_id: str, tool_results: Dict[str, Any]):
        findings = tool_results.get("report", {}).get("detailed_metrics", [])
        if not findings: return

        query = """
        MERGE (e:Episode {session_id: $session_id})
        WITH e
        UNWIND $findings AS finding
        MERGE (r:Resource {name: finding.tool_name})
        SET r.last_observed = datetime(),
            r.current_status = finding.status,
            r.observed_in_session = $session_id

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

        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")
        }
        # Standard Write Handshake
        await self.graph_client._execute_write_query(query, params)

    async def get_infrastructure_ground_truth(self, search_term: str, env: str) -> List[Dict[str, Any]]:
        query = """
        MATCH (ci:ConfigurationItem)-[:DEPENDS_ON*1..2]->(dep)
        WHERE (ci.name CONTAINS $search OR ci.tags CONTAINS $search) AND ci.env = $env
        RETURN ci.name as name, ci.status as status, ci.kind as kind, collect(dep.name) as impacted_services
        LIMIT 25
        """
        # Standard Read Handshake
        return await self.graph_client._execute_read_query(query, {"search": search_term, "env": env})

    # ============================================================================
    # 2. LEELA (Narrative): Episodic History
    # ============================================================================

    async def record_session_narrative(self, session_id: str, results: Dict[str, Any]):
        try:
            agent_specs = await self.tool_registry.load_agentic_capabilities()
            capabilities_snapshot = list(agent_specs.keys())
        except Exception as e:
            logger.warning(f"⚠️ Capability snapshot failure: {e}")
            capabilities_snapshot = []

        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
        WITH e
        MATCH (m:MissionBlueprint {template_id: $template_id})
        MERGE (e)-[:EXECUTED_AS]->(m)
        """
        params = {
            "session_id": session_id,
            "timestamp": datetime.utcnow().isoformat(),
            "outcome": results.get("status", "unknown"),
            "summary": results.get("reflection", "No synthesis provided."),
            "capabilities": capabilities_snapshot,
            "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)

    # ============================================================================
    # 3. DHARMA (Wisdom): Learning & Hindsight
    # ============================================================================

    async def capture_episode(self, session_id: str, context: Dict[str, Any]):
        template_id = context.get("active_template_id", "TPL-GENERIC")
        raw_lessons = context.get("reflection", [])
        if isinstance(raw_lessons, dict): raw_lessons = [raw_lessons]

        for lesson in raw_lessons:
            await self.graph_client.create_lesson_node(
                template_id=template_id,
                topic=lesson.get("topic", "General Improvement"),
                type=lesson.get("type", "OPERATIONAL"),
                related_tool=lesson.get("tool_id"),
                content=lesson.get("insight", "No insight provided."),
                session_id=session_id
            )

    async def get_hindsight(self, current_task: str, template_id: Optional[str] = None) -> List[Dict[str, Any]]:
        query = """
        MATCH (m:MissionBlueprint)-[:GENERATED_LESSON]->(l:ACE_Lesson)
        WHERE (m.template_id = $template_id) OR (toLower(l.topic) CONTAINS toLower($task))
        OPTIONAL MATCH (e:Episode {session_id: l.derived_from})
        RETURN e.outcome AS past_outcome, l.content AS advice, l.topic AS lesson_topic
        ORDER BY l.created_at DESC LIMIT 3
        """
        # 🎯 FIX: Use _execute_read_query for the Read handshake
        return await self.graph_client._execute_read_query(query, {"task": current_task, "template_id": template_id})