"""
Graph Database Client for Neo4j operations.

This module provides functionality for:
- Mission DNA persistence and retrieval
- Context bullet management and scoring
- Infrastructure fact storage and retrieval
- Shared blackboard operations for multi-agent coordination
- Template management with migration support
"""

import datetime
import json
import uuid
from typing import Dict, Any, Optional, List, Literal, Union

import httpx
import yaml

from guardianhub import get_logger
from guardianhub.config.settings import settings
from guardianhub.models.common.common import KeywordList
from guardianhub.models.mission_blueprint import MissionBlueprintDNA
from guardianhub.models.template.suggestion import TemplateSchemaSuggestion

# Module logger
logger = get_logger(__name__)


class GraphDBClient:
    """
    Neo4j Graph Database Client for mission-critical operations.
    
    Provides high-level interface for:
    - Mission Blueprint DNA management
    - Context bullet curation and scoring
    - Infrastructure fact tracking
    - Multi-agent shared blackboard coordination
    """

    # ============================================================================
    # INITIALIZATION & CONFIGURATION
    # ============================================================================

    def __init__(self, poll_interval: int = 5, poll_timeout: int = 300) -> None:
        """Initialize the GraphDB client with configuration from settings."""
        self.api_url = settings.endpoints.get("GRAPH_DB_URL")
        self.headers = {"Accept": "application/json"}
        self.poll_interval = poll_interval
        self.poll_timeout = poll_timeout

        self.client = httpx.AsyncClient(
            headers=self.headers, 
            base_url=self.api_url, 
            timeout=self.poll_timeout + 60
        )
        logger.info(f"🔗 GraphDBClient initialized for URL: {self.api_url}")

    # ============================================================================
    # CORE HELPER METHODS
    # ============================================================================

    async def _execute_write_query(self, query: str, parameters: Dict[str, Any]) -> bool:
        """Execute a write query against the Graph DB service."""
        try:
            response = await self.client.post(
                "/execute-cypher-write",
                json={"query": query, "parameters": parameters},
                timeout=30.0
            )
            response.raise_for_status()
            result = response.json().get("status") == "success"
            if result:
                logger.debug(f"✅ Write query executed successfully")
            return result
        except Exception as e:
            logger.error(f"❌ GraphDB Write Error: {e}")
            return False

    # ============================================================================
    # MISSION DNA MANAGEMENT
    # ============================================================================

    async def save_mission_dna(self, dna: MissionBlueprintDNA) -> bool:
        """
        Anchors the MissionBlueprintDNA into the graph.
        Uses MERGE to ensure stability and updates properties if evolved.
        """
        query = """
        MERGE (m:MissionBlueprint {template_id: $template_id})
        SET m.mission_category = $mission_category,
            m.brief_summary = $brief_summary,
            m.mission_rationale = $mission_rationale,
            m.target_persona = $target_persona,
            m.auth_level_required = $auth_level_required,
            m.success_criteria = $success_criteria,
            m.safety_constraints = $safety_constraints,
            m.escalation_policy = $escalation_policy,
            m.estimated_blast_radius = $estimated_blast_radius,
            m.execution_timeout_seconds = $execution_timeout_seconds,
            m.briefing_template = $briefing_template,
            m.version = $version,
            m.last_evolved_at = datetime(),
            m.reflection_count = coalesce(m.reflection_count, 0) + $reflection_inc

        // Link to the overarching Category node for topological discovery
        MERGE (c:MissionCategory {name: $mission_category})
        MERGE (m)-[:BELONGS_TO]->(c)

        RETURN m.template_id as id
        """

        params = {
            "template_id": dna.template_id,
            "mission_category": dna.mission_category,
            "brief_summary": dna.brief_summary,
            "mission_rationale": dna.mission_rationale,
            "target_persona": dna.target_persona,
            "auth_level_required": dna.auth_level_required,
            "success_criteria": dna.success_criteria,  # Neo4j handles lists of strings
            "safety_constraints": dna.safety_constraints,
            "escalation_policy": dna.escalation_policy,
            "estimated_blast_radius": dna.estimated_blast_radius,
            "execution_timeout_seconds": dna.execution_timeout_seconds,
            "briefing_template": dna.briefing_template,
            "version": dna.version,
            "reflection_inc": 1 if dna.reflection_count > 0 else 0
        }

        logger.info(f"⚓ Anchoring Mission DNA: {dna.template_id} [{dna.mission_category}]")
        result = await self._execute_write_query(query, params)
        if result:
            logger.info(f"✅ Mission DNA successfully anchored: {dna.template_id}")
        else:
            logger.error(f"❌ Failed to anchor Mission DNA: {dna.template_id}")
        return result

    # ============================================================================
    # TEMPLATE & DNA RETRIEVAL
    # ============================================================================

    async def get_template_by_id(self, template_id: str) -> Optional[Dict[str, Any]]:
        """
        Retrieves MissionBlueprint properties.
        Supports legacy DocumentTemplate lookup for migration safety.
        """
        query = """
        MATCH (m) 
        WHERE (m:MissionBlueprint OR m:DocumentTemplate) AND m.template_id = $template_id
        RETURN properties(m) as template, labels(m)[0] as type
        """
        try:
            response = await self.client.post(
                "/query-cypher",
                json={"query": query, "parameters": {"template_id": template_id}},
                timeout=30.0
            )
            response.raise_for_status()
            data = response.json()

            if data.get("status") == "success" and data.get("results"):
                logger.debug(f"✅ Retrieved template/DNA: {template_id}")
                return data["results"][0]["template"]
            
            logger.warning(f"⚠️ Template/DNA not found: {template_id}")
            return None
        except Exception as e:
            logger.error(f"❌ Failed to retrieve DNA: {template_id} - {e}")
            return None

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

    async def merge_context_bullet(self, template_id: str, bullet_data: Dict[str, Any]) -> bool:
        """
        ACE Curator: Creates a new :ContextBullet node and links it.
        (UPDATED to include domain, keywords, and curator)
        """
        bullet_id = bullet_data.get("bullet_id", f"BUL-{uuid.uuid4().hex[:8]}")

        query = """
            // 🟢 FINAL FIX: MERGE the template node first to guarantee its existence.
            MERGE (t:DocumentTemplate {template_id: $template_id})
            ON CREATE SET t.created_at = datetime() 

            // MERGE or create the ContextBullet node
            MERGE (b:ContextBullet {bullet_id: $bullet_id})
            ON CREATE SET 
                b.content = $content,
                b.type = $type,
                b.domain = $domain,
                b.source_curator = $source_curator,
                b.helpful_count = 0,
                b.harmful_count = 0,
                b.created_at = datetime()

            // MERGE the relationship
            MERGE (t)-[:HAS_BULLET]->(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"),
            # --- NEW PARAMETERS ---
            "domain": bullet_data.get("domain", "DEFAULT"),
            "keywords": bullet_data.get("keywords", []),
            "source_curator": bullet_data.get("source_curator", "system-ingest")
        }

        logger.info(f"🎯 Curating new bullet {bullet_id} for template {template_id}")
        result = await self._execute_write_query(query, params)
        if result:
            logger.info(f"✅ Successfully curated bullet: {bullet_id}")
        else:
            logger.error(f"❌ Failed to curate bullet: {bullet_id}")
        return result

    async def update_bullet_scores(self, bullet_id: str, feedback_type: Literal['helpful', 'harmful']) -> bool:
        """Update bullet scores based on user feedback."""
        logger.debug(f"📊 Updating bullet {bullet_id} with feedback: {feedback_type}")
        
        # Determine the score increment clause
        if feedback_type == 'helpful':
            score_clause = "b.helpful_count = coalesce(b.helpful_count, 0) + 1"
        elif feedback_type == 'harmful':
            score_clause = "b.harmful_count = coalesce(b.harmful_count, 0) + 1"
        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() 
        RETURN b
        """

        result = await self._execute_write_query(query, {"bullet_id": bullet_id})
        if result:
            logger.debug(f"✅ Updated bullet score: {bullet_id} ({feedback_type})")
        else:
            logger.error(f"❌ Failed to update bullet score: {bullet_id}")
        return result

    async def get_top_bullets_for_template(
            self,
            template_id: str,
            user_persona: str,
            limit: int = 5
    ) -> List[Dict[str, Any]]:
        """
        Retrieves top strategies, prioritizing those aligned with the user's role.
        """
        logger.debug(f"🔍 Retrieving top bullets for {template_id} (persona: {user_persona})")
        
        query = """
            MATCH (m:MissionBlueprint {template_id: $template_id})-[:HAS_BULLET]->(b:ContextBullet)

            WITH b,
                 (coalesce(b.helpful_count, 0) - coalesce(b.harmful_count, 0)) AS utility,
                 CASE WHEN b.target_persona = $user_persona THEN 1.2 ELSE 1.0 END AS persona_multiplier,
                 toFloat(duration.between(coalesce(b.created_at, datetime()), datetime()).days) AS days_old

            WITH b, (utility * persona_multiplier) / (1.0 + (days_old / 30.0)) AS final_score
            RETURN properties(b) AS bullet, final_score
            ORDER BY final_score DESC
            LIMIT $limit
            """

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

        try:
            response = await self.client.post("/query-cypher", json=payload, timeout=30.0)
            response.raise_for_status()
            results = response.json().get("results", [])
            
            logger.debug(f"📊 Retrieved {len(results)} bullets for {template_id}")

            # Format strictly as expected by the Assembly Engine
            return [
                {
                    "id": rec["bullet"]["bullet_id"],
                    "content": rec["bullet"]["content"],
                    "type": rec["bullet"].get("type", "STRATEGY"),
                    "score": 1.0 + (rec["bullet"].get("helpful_count", 0) * 0.1)
                }
                for rec in results
            ]
        except Exception as e:
            logger.error(f"❌ Failed to fetch top bullets for {template_id}: {e}")
            return []

    async def get_bullets_by_id(self, bullet_ids: List[str]) -> List[Dict[str, Any]]:
        """
        ACE Generator: Retrieves bullet metadata (counts/type) for the ContextAssemblyEngine.
        """
        if not bullet_ids:
            return []

        query = """
        MATCH (b:ContextBullet)
        WHERE b.bullet_id IN $bullet_ids
        RETURN properties(b) as bullet
        """

        try:
            response = await self.client.post(
                "/query-cypher",
                json={"query": query, "parameters": {"bullet_ids": bullet_ids}},
                timeout=30.0
            )
            response.raise_for_status()
            data = response.json()

            if data.get("status") == "success":
                return [rec['bullet'] for rec in data.get('results', [])]
            return []
        except Exception as e:
            logger.error(f"Failed to retrieve bullets: {e}")
            return []

    async def get_infrastructure_facts(
            self,
            template_id: str,
            keywords: Optional[List[str]] = None,
    ) -> List[Dict[str, Any]]:
        """
        Retrieve topological 'Ground Truth' anchored to the Mission Pattern.
        """
        if not keywords:
            logger.warning(f"No keywords for fact retrieval in {template_id}. Returning empty set.")
            return []

        # 1. THE ANCHORED SEARCH
        # We use a pattern-based MATCH instead of OPTIONAL to ensure we only get relevant data.
        # We support both label types (Blueprint/Template) for migration safety.
        cypher_query = """
        MATCH (m) 
        WHERE (m:MissionBlueprint OR m:DocumentTemplate) AND m.template_id = $template_id

        // We only care about facts that actually exist for this mission
        MATCH (m)-[:HAS_RECENT_FACT]->(f:InfrastructureFact)
        OPTIONAL MATCH (r:Resource)-[:PROVIDES_ASSERTION]->(f)

        // Perform Case-Insensitive topological filtering
        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,
            updated: toString(f.updated_at) 
        } as fact
        ORDER BY fact.updated DESC
        LIMIT 5
        """

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

        try:
            response = await self.client.post(
                "/query-cypher",
                json={"query": cypher_query, "parameters": params},
                timeout=30.0
            )
            response.raise_for_status()
            results = response.json().get('results', [])

            enriched_facts = []
            for record in results:
                fact = record.get("fact", {})
                # 🧠 Safe JSON hydration
                if fact.get("metrics_raw"):
                    try:
                        fact["metrics"] = json.loads(fact["metrics_raw"])
                    except (json.JSONDecodeError, TypeError):
                        fact["metrics"] = {}
                    del fact["metrics_raw"]
                enriched_facts.append(fact)

            logger.info(f"📊 Retrieved {len(enriched_facts)} infrastructure facts for {template_id}")
            return enriched_facts

        except Exception as e:
            logger.error(f"❌ Fact 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]]:
        """
        Retrieve ACE lessons based on query or pre-extracted keywords.

        Args:
            query: The search query string (optional if keywords are provided)
            keywords: Pre-extracted keywords or KeywordList object
            limit: Maximum number of results to return

        Returns:
            List of lesson dictionaries with id, content, type, and related_tool
        """
        # Convert KeywordList to list of strings if needed
        if not keywords:
            logger.warning("No valid keywords provided for ACE lessons search")
            return []

        logger.debug(f"Searching ACE lessons with keywords: {keywords}")

        cypher_query = """
        MATCH (b:ContextBullet)
        WHERE any(kw IN $keywords WHERE 
            toLower(b.content) CONTAINS toLower(kw) OR 
            toLower(b.domain) CONTAINS toLower(kw)
        )
        RETURN {
            id: b.bullet_id,
            content: b.content,
            type: b.type,
            related_tool: b.related_tool,
            score: size([kw IN $keywords WHERE 
                toLower(b.content) CONTAINS toLower(kw) OR 
                toLower(b.domain) CONTAINS toLower(kw)
            ]) * 1.0 / size($keywords)
        } as lesson
        ORDER BY lesson.score DESC
        LIMIT $limit
        """

        try:
            response = await self.client.post(
                "/query-cypher",
                json={
                    "query": cypher_query,
                    "parameters": {
                        "keywords": keywords,
                        "limit": limit
                    }
                },
                timeout=30.0
            )
            response.raise_for_status()

            results = response.json().get("results", [])
            return [record["lesson"] for record in results if "lesson" in record]

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

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

    async def upsert_shared_resources(
            self,
            session_id: str,
            findings: List[Dict[str, Any]],
            template_id: str
    ) -> bool:
        """
        Promotes Specialist findings to the Shared Blackboard.
        Anchors facts to the MissionBlueprint to prevent 'Intelligence Drift'.
        """
        if not findings: return True

        # Pre-serialize metrics for Cypher compatibility
        for f in findings:
            f["metrics_json"] = json.dumps(f.get("metrics", {}))

        query = """
        UNWIND $findings AS finding
        // Anchor to the Sovereign Pattern
        MATCH (m:MissionBlueprint {template_id: $active_template_id})

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

        // Create the Immutable Fact for this specific Mission Instance
        MERGE (f:InfrastructureFact {fact_id: "F-" + $active_template_id + "-" + finding.tool_name})
        SET f.content = finding.factual_summary_text,
            f.metrics = finding.metrics_json,
            f.updated_at = datetime()

        // Topological Binding
        MERGE (r)-[:PROVIDES_ASSERTION]->(f)
        MERGE (m)-[:HAS_RECENT_FACT]->(f)

        RETURN count(f) as fact_count
        """
        return await self._execute_write_query(query, {
            "session_id": session_id,
            "findings": findings,
            "active_template_id": template_id
        })


    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._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
