import uuid
from datetime import datetime
import json
from typing import Dict, Any, List, Optional, Literal
from guardianhub import get_logger
from guardianhub.clients import VectorClient, GraphDBClient, ToolRegistryClient
# 🎯 Standardized Contracts
from guardianhub.models import VectorQueryRequest, VectorQueryResponse

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]):
        """Promotes tool findings to Graph Facts (Satya)."""
        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
        }
        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]]:
        """
        Avatar: Kurma (The Shell/Foundation).
        Retrieves CI topology from Neo4j Satya Segment.
        """
        logger.info(f"🐢 [SATYA] Querying topology for: {search_term} ({env})")

        # Clean the search term to prevent Cypher injection or empty matches
        clean_term = search_term.replace("'", "").replace("\"", "")

        query = """
        MATCH (ci:ConfigurationItem)
        WHERE (ci.name CONTAINS $search OR any(tag IN ci.tags WHERE tag CONTAINS $search))
        AND ci.env = $env
        OPTIONAL MATCH (ci)-[:DEPENDS_ON*1..2]->(dep)
        RETURN ci.name as name, ci.status as status, ci.kind as kind, 
               collect(DISTINCT dep.name) as impacted_services
        LIMIT 25
        """
        try:
            results = await self.graph_client._execute_read_query(
                query,
                {"search": clean_term, "env": env}
            )
            return results if results else []
        except Exception as e:
            logger.error(f"❌ [SATYA] Graph fetch failed: {e}")
            return []  # Fail-open with empty list to satisfy workflow types

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

    async def record_session_narrative(self, session_id: str, results: Dict[str, Any]):
        """Records the mission story in the Graph and caches it in Vector Memory."""
        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 = []

        # 1. Update Graph Narrative
        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)

        # 🎯 2. Standardized Vector Upsert
        # We push a stringified narrative to the episodes collection
        await self.vector_client.upsert_atomic(
            collection="episodes",
            doc_id=f"narrative-{session_id}",
            text=results.get("reflection", "Mission completion without summary."),
            metadata={
                "session_id": session_id,
                "template_id": params["template_id"],
                "type": "mission_narrative",
                "outcome": params["outcome"]
            }
        )

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

    async def get_recent_episodes(self, query_text: str, template_id: str) -> List[Dict[str, Any]]:
        """
        [HINDSIGHT HANDSHAKE]
        Uses the Standardized Vector Contract to retrieve historical parallels.
        """
        try:
            # 🚀 THE 10-FOLD MOVE: No more dict interpreting.
            request = VectorQueryRequest(
                query=query_text,
                collection="episodes",
                n_results=5,
                filters={"template_id": template_id}
            )

            response: VectorQueryResponse = await self.vector_client.query(request)

            # Return validated content only
            return [res.model_dump() for res in response.results]
        except Exception as e:
            logger.error(f"❌ [HINDSIGHT] Search failed: {e}")
            return []

    async def get_hindsight(self, current_task: str, template_id: Optional[str] = None) -> List[Dict[str, Any]]:
        """Cross-references Graph Lessons with Episode outcomes."""
        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
        """
        return await self.graph_client._execute_read_query(query, {"task": current_task, "template_id": template_id})

        # ============================================================================
        # CONTEXT BULLET MANAGEMENT
        # ============================================================================

    async def merge_context_bullet(self, template_id: str, bullet_data: Dict[str, Any]) -> bool:
        """
        [DHARMA EVOLUTION]
        ACE Curator: Anchors a new Strategic Directive to the Mission DNA.
        Converts a Specialist's 'Reflection' into the system's 'Operating Law'.
        """
        bullet_id = bullet_data.get("bullet_id", f"BUL-{uuid.uuid4().hex[:8]}")

        query = """
            // 1. MISSION DNA ANCHOR: Ensure the Blueprint exists
            MERGE (m:MissionBlueprint {template_id: $template_id})
            ON CREATE SET m.created_at = datetime() 

            // 2. WISDOM ANCHOR: Create/Update the Directive
            MERGE (b:ContextBullet {bullet_id: $bullet_id})
            ON CREATE SET 
                b.content = $content,
                b.type = $type,
                b.domain = $domain,
                b.keywords = $keywords,
                b.source_curator = $source_curator,
                b.helpful_count = 1,  // Starts with 1 because it was just proven
                b.harmful_count = 0,
                b.created_at = datetime()
            ON MATCH SET
                b.keywords = $keywords, // Keep keywords fresh
                b.last_proven_at = datetime()

            // 3. TOPOLOGICAL BINDING: Use the 'Evolved' relationship
            MERGE (m)-[:EVOLVED_STRATEGY]->(b)

            RETURN count(b) AS merge_count
            """

        params = {
            "template_id": template_id,
            "bullet_id": bullet_id,
            "content": bullet_data.get("content", ""),
            "type": bullet_data.get("type", "STRATEGY"),
            "domain": bullet_data.get("domain", "Sovereign"),
            "keywords": bullet_data.get("keywords", []),
            "source_curator": bullet_data.get("source_curator", "Specialist-Reflection")
        }

        logger.info(f"🎯 [CURATOR] Anchoring Wisdom {bullet_id} to Blueprint {template_id}")
        result = await self.graph_client._execute_write_query(query, params)

        if result:
            logger.info(f"✅ [CURATOR] Strategy evolved: {bullet_id}")
        else:
            logger.error(f"❌ [CURATOR] Failed to anchor wisdom: {bullet_id}")
        return result

    async def update_bullet_scores(
            self,
            bullet_id: str,
            feedback_type: Literal['helpful', 'harmful'],
            weight: float = 1.0  # 🚀 New: Allow weighted feedback
    ) -> bool:
        """
        [KARMA TUNING]
        Updates the utility metrics of a Strategic Directive.
        Used to prune ineffective strategies and boost proven wisdom.
        """
        logger.debug(f"📊 [FEEDBACK] Tuning Wisdom {bullet_id}: {feedback_type} (Weight: {weight})")

        # 🎯 Penalty Bias: In Sovereign systems, 'Harmful' acts have higher weight by default
        effective_weight = weight if feedback_type == 'helpful' else weight * 2.0

        if feedback_type == 'helpful':
            score_clause = "b.helpful_count = coalesce(b.helpful_count, 0) + $weight"
        elif feedback_type == 'harmful':
            score_clause = "b.harmful_count = coalesce(b.harmful_count, 0) + $weight"
        else:
            logger.warning(f"⚠️ Invalid feedback type: {feedback_type}")
            return False

        query = f"""
        MATCH (b:ContextBullet {{bullet_id: $bullet_id}}) 
        SET {score_clause}, 
            b.last_feedback_at = datetime(),
            b.total_interactions = coalesce(b.total_interactions, 0) + 1
        RETURN b.bullet_id as id
        """

        params = {"bullet_id": bullet_id, "weight": effective_weight}

        result = await self.graph_client._execute_write_query(query, params)

        if result:
            logger.info(f"✅ [FEEDBACK] Wisdom calibrated: {bullet_id} increased {feedback_type}")
        return result

    async def get_top_directives_for_mission(
            self,
            template_id: str,
            user_persona: str,
            limit: int = 5
    ) -> List[Dict[str, Any]]:
        """
        [DHARMA RETRIEVAL]
        Retrieves the most effective tactical directives anchored to a Mission DNA.
        Prioritizes proven strategies that align with the Blueprint's history.
        """
        logger.debug(f"🔍 Harvesting Top Directives for DNA: {template_id}")

        query = """
            MATCH (m:MissionBlueprint {template_id: $template_id})-[:EVOLVED_STRATEGY]->(b:ContextBullet)

            // 1. Calculate the 'Wisdom Score'
            WITH b,
                 (coalesce(b.helpful_count, 0) - (coalesce(b.harmful_count, 0) * 2.0)) AS net_utility,
                 CASE WHEN b.target_persona = $user_persona THEN 1.5 ELSE 1.0 END AS persona_weight,
                 // Decay score for stale strategies (Freshness factor)
                 toFloat(duration.between(coalesce(b.created_at, datetime()), datetime()).days) AS days_old

            // 2. Apply Sovereign Weighting: utility * persona / decay
            WITH b, (net_utility * persona_weight) / (1.0 + (days_old / 45.0)) AS wisdom_score

            RETURN properties(b) AS directive, wisdom_score
            ORDER BY wisdom_score DESC
            LIMIT $limit
            """

        payload = {
                "template_id": template_id,
                "user_persona": user_persona,
                "limit": limit
            }

        try:
            results = await self.graph_client._execute_read_query(query, parameters=payload)

            # 🎯 The "10-fold" Return: Map directives back to the Assembly Engine format
            directives = []
            for rec in results:
                d = rec["directive"]
                directives.append({
                    "id": d["bullet_id"],
                    "content": d["content"],
                    "type": d.get("type", "STRATEGY"),
                    "domain": d.get("domain", "Sovereign"),
                    # We normalize the internal wisdom_score for the LLM's context scoring
                    "score": round(rec["wisdom_score"], 2)
                })

            logger.info(f"🧠 [WISDOM_FETCH] Loaded {len(directives)} directives for DNA {template_id}")
            return directives

        except Exception as e:
            logger.error(f"❌ [WISDOM_FETCH] Failed to harvest directives for {template_id}: {e}")
            return []

    async def get_infrastructure_facts(
            self,
            template_id: str,
            keywords: Optional[List[str]] = None,
            cross_template: bool = False  # 🚀 New: The Intelligence Handshake
    ) -> List[Dict[str, Any]]:
        """
        [SATYA RETRIEVAL]
        Retrieves topological 'Ground Truth'.
        If cross_template=True, it allows 'overhearing' facts from other agent domains.
        """
        if not keywords:
            logger.warning(f"⚠️ No keywords for fact retrieval in {template_id}.")
            return []

        # 🎯 THE EVOLVED QUERY
        # If cross_template, we drop the specific MissionBlueprint anchor and look at Resources directly.
        anchor_clause = "MATCH (m:MissionBlueprint {template_id: $template_id})-[:HAS_RECENT_FACT]->(f:InfrastructureFact)"
        if cross_template:
            anchor_clause = "MATCH (f:InfrastructureFact)"

        cypher_query = f"""
        {anchor_clause}
        OPTIONAL MATCH (r:Resource)-[:PROVIDES_ASSERTION]->(f)

        WITH f, r, [kw IN $keywords | toLower(kw)] AS lower_kws
        WHERE any(kw IN lower_kws WHERE 
            toLower(f.content) CONTAINS kw OR 
            toLower(r.name) CONTAINS kw
        )

        RETURN DISTINCT {{
            resource: coalesce(r.name, "system"),
            summary: f.content,
            metrics_raw: f.metrics,
            observed_at: toString(f.updated_at),
            source_agent: f.source_agent
        }} as fact
        ORDER BY fact.observed_at DESC
        LIMIT 10
        """

        params = {"template_id": template_id, "keywords": keywords}

        try:
            results = await self.graph_client._execute_read_query(
                cypher_query,
                params,
            )

            # 🧠 HYDRATION LAYER: Reconstructing the sensory data
            enriched_facts = []
            for record in results:
                fact = record.get("fact", {})
                if fact.get("metrics_raw"):
                    try:
                        # Metrics were stored as JSON strings in Neo4j to preserve types
                        fact["metrics"] = json.loads(fact["metrics_raw"])
                    except (json.JSONDecodeError, TypeError):
                        fact["metrics"] = {}
                    del fact["metrics_raw"]
                enriched_facts.append(fact)

            logger.info(f"📊 [SATYA_LOAD] Inherited {len(enriched_facts)} infrastructure facts for {template_id}")
            return enriched_facts

        except Exception as e:
            logger.error(f"❌ [SATYA_LOAD] Retrieval failed for {template_id}: {e}", exc_info=True)
            return []


    async def get_ace_lessons(
            self,
            query: Optional[str] = None,
            keywords: Optional[List[str]] = None,
            limit: int = 10
    ) -> List[Dict[str, Any]]:
        """
        [DHARMA RETRIEVAL]
        Retrieve ACE Lessons (Beliefs) from the Agent Space.
        Prioritizes lessons learned by similar Specialists or linked to active Tools.
        """
        if not keywords:
            logger.warning("⚠️ [WISDOM] No keywords provided. Agent is flying blind.")
            return []

        # 1. THE RELATIONAL SEARCH
        # We look for lessons in the ContextBullet pool that match keywords,
        # but we also check if they are linked to a successful Episode.
        cypher_query = """
        MATCH (b:ContextBullet)
        WHERE any(kw IN $keywords WHERE 
            toLower(b.content) CONTAINS toLower(kw) OR 
            toLower(b.domain) CONTAINS toLower(kw)
        )

        // Optional: Tie back to the Specialist or Blueprint for 'Success Proof'
        OPTIONAL MATCH (m:MissionBlueprint)-[r:EVOLVED_STRATEGY]->(b)

        WITH b, count(r) as proven_count, [kw IN $keywords | toLower(kw)] as lower_kws

        // Calculate a cognitive score based on keyword density + historical proof
        WITH b, proven_count,
             size([kw IN lower_kws WHERE toLower(b.content) CONTAINS kw]) as match_density

        RETURN {
            id: b.bullet_id,
            content: b.content,
            type: b.type,
            domain: b.domain,
            related_tool: b.related_tool,
            proven_utility: proven_count,
            score: (match_density * 1.0 / size($keywords)) + (proven_count * 0.1)
        } as lesson
        ORDER BY lesson.score DESC
        LIMIT $limit
        """

        try:
            results = await self.graph_client._execute_read_query(
                cypher_query,{"keywords": keywords, "limit": limit}
            )

            lessons = [record["lesson"] for record in results if "lesson" in record]
            logger.info(f"🧠 [WISDOM_LOAD] Harvested {len(lessons)} actionable lessons from Agent Space.")
            return lessons

        except Exception as e:
            logger.error(f"❌ [WISDOM_LOAD] Failed to retrieve agentic wisdom: {e}", exc_info=True)
            return []

    # ==========================================================================
    # 2. SHARED BLACKBOARD (Mission-Anchored)
    # ==========================================================================

    async def anchor_mission_findings(
            self,
            session_id: str,
            template_id: str,
            findings: List[Dict[str, Any]],
            source_agent: str = "specialist"
    ) -> bool:
        """
        [SATYA CONSOLIDATION]
        The Single Source of Truth for promoting field discoveries.
        Links the 'Fact' to the 'Mission DNA' and the 'Episode'.
        """
        if not findings:
            return True

        # Prepare findings with serialized metrics
        processed_findings = []
        for f in findings:
            processed_findings.append({
                **f,
                "fact_id": f"F-{template_id}-{f.get('tool_name', 'unknown').replace(' ', '_')}",
                "metrics_json": json.dumps(f.get("metrics", {}))
            })

        query = """
        UNWIND $findings AS finding

        // 1. ANCHOR: Locate the Mission DNA and the Session Episode
        MATCH (m:MissionBlueprint {template_id: $template_id})
        MERGE (e:Episode {session_id: $session_id})

        // 2. RESOURCE: MERGE the global CI/Resource
        MERGE (r:Resource {name: finding.tool_name})
        SET r.last_observed = datetime(),
            r.current_status = finding.status

        // 3. FACT: Create/Update the Immutable Fact
        MERGE (f:InfrastructureFact {fact_id: finding.fact_id})
        SET f.content = finding.factual_summary_text,
            f.metrics = finding.metrics_json,
            f.updated_at = datetime(),
            f.source_agent = $source_agent

        // 4. TOPOLOGY: Establish the Triple-Bind
        MERGE (m)-[:HAS_RECENT_FACT]->(f)  // DNA Link
        MERGE (e)-[:DISCOVERED]->(f)        // Session Link
        MERGE (r)-[:PROVIDES_ASSERTION]->(f) // Resource Link

        RETURN count(f) as fact_count
        """

        params = {
            "session_id": session_id,
            "template_id": template_id,
            "findings": processed_findings,
            "source_agent": source_agent
        }

        logger.info(f"⚓ [CONSOLIDATOR] Anchoring {len(findings)} findings for session {session_id}")
        return await self.graph_client._execute_write_query(query, params)

    async def upsert_infrastructure_fact(
            self,
            template_id: str,
            resource_name: str,
            fact_content: str,
            metrics: Optional[Dict[str, Any]] = None,
            source_agent: str = "system"
    ) -> bool:
        """
        SENSORY INPUT: Records a new finding into the graph.
        Anchors the Fact to the Mission Pattern and the specific Resource.
        """

        # 1. Prepare the Fact ID to prevent duplication within a mission context
        # format: F-<TEMPLATE>-<RESOURCE_NAME>
        fact_id = f"F-{template_id}-{resource_name.replace(' ', '_')}"

        query = """
        // 1. Ensure the Mission anchor exists
        MATCH (m) 
        WHERE (m:MissionBlueprint OR m:DocumentTemplate) AND m.template_id = $template_id

        // 2. MERGE the Resource (The 'What')
        MERGE (r:Resource {name: $resource_name})
        SET r.last_observed = datetime(),
            r.observed_by = $source_agent

        // 3. MERGE the Fact (The 'Intelligence')
        MERGE (f:InfrastructureFact {fact_id: $fact_id})
        SET f.content = $content,
            f.metrics = $metrics_json,
            f.updated_at = datetime(),
            f.source_agent = $source_agent

        // 4. Establish Topological Links
        MERGE (r)-[:PROVIDES_ASSERTION]->(f)
        MERGE (m)-[:HAS_RECENT_FACT]->(f)

        RETURN f.fact_id as id
        """

        params = {
            "template_id": template_id,
            "resource_name": resource_name,
            "fact_id": fact_id,
            "content": fact_content,
            "metrics_json": json.dumps(metrics or {}),
            "source_agent": source_agent
        }

        try:
            success = await self.graph_client._execute_write_query(query, params)
            if success:
                logger.info(f"🛰️ Fact promoted to Graph: {fact_id} (Source: {source_agent})")
            return success
        except Exception as e:
            logger.error(f"❌ Failed to promote fact {fact_id}: {e}")
            return False

    async def create_lesson_node(
            self,
            template_id: str,
            topic: str,
            type: str,
            content: str,
            related_tool: Optional[str] = None,
            session_id: Optional[str] = None
    ) -> bool:
        """
        [DHARMA ANCHOR]
        Creates an ACE_Lesson node and anchors it to the MissionBlueprint.
        This allows the system to 'learn' from this specific mission type.
        """
        query = """
           // 1. Ensure the Mission DNA exists as the anchor
           MATCH (m:MissionBlueprint {template_id: $template_id})

           // 2. MERGE the Lesson based on content to prevent duplicates
           MERGE (l:ACE_Lesson {content: $content})
           SET l.topic = $topic,
               l.type = $type,
               l.related_tool = $related_tool,
               l.derived_from = $session_id,
               l.created_at = datetime()

           // 3. Link the Wisdom to the Law (The DNA)
           MERGE (m)-[:GENERATED_LESSON]->(l)

           RETURN count(l) as lesson_count
           """
        params = {
            "template_id": template_id,
            "topic": topic,
            "type": type,
            "content": content,
            "related_tool": related_tool,
            "session_id": session_id
        }
        return await self.graph_client._execute_write_query(query, params)


    async def get_bullets_by_id(self, bullet_ids: List[str]) -> List[Dict[str, Any]]:
        """
        [DHARMA VALIDATION]
        Retrieves detailed metadata for specific wisdom bullets.
        Used by the ContextAssemblyEngine to perform final 'Truth Scoring'.
        """
        if not bullet_ids:
            return []

        # We fetch the properties + the count of how many Blueprints currently trust this bullet
        query = """
        MATCH (b:ContextBullet)
        WHERE b.bullet_id IN $bullet_ids

        OPTIONAL MATCH (m:MissionBlueprint)-[r:EVOLVED_STRATEGY]->(b)

        RETURN properties(b) as properties, 
               count(m) as trust_factor,
               labels(b) as labels
        """

        try:
            results = await self.graph_client._execute_read_query(
                query,
                {"bullet_ids": bullet_ids}
            )
            enriched_bullets = []
            for rec in results:
                bullet = rec['properties']
                # Inject structural metadata for the Assembly Engine
                bullet["trust_factor"] = rec["trust_factor"]
                bullet["is_ace_bullet"] = "ContextBullet" in rec["labels"]
                enriched_bullets.append(bullet)

                logger.debug(f"📊 [WISDOM_LOAD] Re-hydrated {len(enriched_bullets)} bullets from ID list.")
                return enriched_bullets

            return []
        except Exception as e:
            logger.error(f"❌ [WISDOM_LOAD] Failed to retrieve wisdom lineage: {e}")
            return []
