from datetime import datetime
from typing import Dict, Any, List


class EpisodicManager:
    def __init__(self, vector_client, graph_client, tool_registry):
        self.vector_client = vector_client
        self.graph_client = graph_client
        self.tool_registry = tool_registry

    async def record_resource_discovery(self, session_id: str, tool_results: Dict[str, Any]):
        """
        Promotes tool findings to permanent 'Shared Resource' nodes in the Graph.
        This ensures the NEXT session starts with these 'Facts'.
        """
        # 1. Extract findings (e.g., the 42 issues from your logs)
        findings = tool_results.get("report", {}).get("detailed_metrics", [])

        cypher_query = """
        UNWIND $findings AS finding
        // Create/Update the Resource (The Common Language)
        MERGE (r:Resource {resource_id: finding.tool_name + "_" + finding.status})
        SET r.name = finding.tool_name,
            r.last_observed = datetime(),
            r.current_status = finding.status,
            r.observed_in_session = $session_id,
            r.otel_trace_id = $trace_id

        // Link to the Episode for provenance
        WITH r
        MATCH (e:Episode {session_id: $session_id})
        MERGE (e)-[:DISCOVERED]->(r)

        // Create a 'Fact' assertion that other agents can read
        MERGE (f:Fact {fact_id: "assertion_" + r.resource_id})
        SET f.content = finding.factual_summary_text,
            f.type = 'infrastructure',
            f.last_updated = datetime()
        MERGE (r)-[:HAS_ASSERTION]->(f)
        """

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

        await self.graph_client.query(cypher_query, params)

    async def record_session_narrative(self, session_id: str, results: Dict[str, Any]):
        """Persists the session as an Episode node with its tool snapshot."""
        # 1. Get current capabilities at the MOMENT of completion
        active_tools = await self.tool_registry.get_active_capabilities()
        capabilities = [t['name'] for t in active_tools]

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

        WITH e
        UNWIND $affected_cis AS ci_name
        MATCH (i:InfrastructureEntity {name: ci_name})
        MERGE (e)-[:ENCOUNTERED]->(i)
        """

        params = {
            "session_id": session_id,
            "timestamp": datetime.utcnow().isoformat(),
            "outcome": results.get("status", "unknown"),
            "capabilities": capabilities,  # Snapshot of what was possible
            "affected_cis": results.get("affected_cis", [])
        }
        await self.graph_client.query(query, params)

    async def capture_episode(self, session_id: str, context: Dict[str, Any]):
        # Extract lessons from the LLM reflection
        raw_lessons = context.get("reflection", [])

        for lesson in raw_lessons:
            # Determine if we HAVE the tools to act on this
            if self._is_actionable_in_current_ecosystem(lesson):
                await self.graph_client.create_lesson_node(
                    topic=lesson.topic,
                    type=self._classify_type(lesson),  # Operational, Structural, etc.
                    related_tool=lesson.tool_id
                )

    async def get_hindsight(self, current_task: str):
        """
        Retrieves past episodes to help the agent distinguish
        between its static beliefs and recent outcomes.
        """
        return await self.graph_client.query(
            "MATCH (e:Episode)-[:GENERATED]->(l:ACE_Lesson) "
            "WHERE l.topic CONTAINS $task "
            "RETURN e.outcome, l.suggested_improvement "
            "ORDER BY e.timestamp DESC LIMIT 3"
        )