# services/clients/neo4j_client.py

import datetime
import json
import uuid  # Needed for bullet ID generation
from typing import Dict, Any, Optional, List, Literal, Union
from guardianhub.models.common.common import KeywordList
import httpx
import yaml
from guardianhub import get_logger
from guardianhub.models.template.suggestion import TemplateSchemaSuggestion

from guardianhub.config.settings import settings

logger = get_logger(__name__)


class GraphDBClient:

    def __init__(self, poll_interval: int = 5, poll_timeout: int = 300):
        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("GraphDBClient initialized for URL: %s", self.api_url)

    # --- 1. Template Registration (Existing YAML Logic) ---

    async def save_document_template(self, template: TemplateSchemaSuggestion) -> bool:
        """
        Creates a new DocumentTemplate node via YAML ingestion.
        """
        template_properties = {
            "template_id": template.template_id,
            "document_type": template.document_type,
            "template_name": template.template_name,
            "required_keywords": template.required_keywords,
            "json_schema_str": json.dumps(template.json_schema)
        }

        ingestion_payload = {
            "nodes": [{"type": "DocumentTemplate", "properties": template_properties}],
            "relationships": [
                {
                    "from": {"type": "PlatformService", "property": "name", "value": "doc-template-service"},
                    "to": {"type": "DocumentTemplate", "property": "template_id", "value": template.template_id},
                    "type": "MANAGES_TEMPLATE",
                    "properties": {"link_date": datetime.datetime.now().isoformat()}
                }
            ]
        }

        try:
            yaml_payload = yaml.dump(ingestion_payload, sort_keys=False)
            response = await self.client.post(
                "/ingest-yaml-schema",
                content=yaml_payload,
                headers={'Content-Type': 'application/x-yaml'},
                timeout=30.0
            )
            response.raise_for_status()
            if response.json().get("status") == "success":
                logger.info(f"✅ Template {template.template_id} successfully persisted.")
                return True
            return False
        except Exception as e:
            logger.error(f"Failed to ingest template YAML: {e}", exc_info=True)
            return False

    # --- 2. Standard Retrieval (Existing Logic) ---

    async def get_template_by_id(self, template_id: str) -> Optional[Dict[str, Any]]:
        """Retrieves DocumentTemplate properties by ID."""
        query = "MATCH (t:DocumentTemplate {template_id: $template_id}) RETURN properties(t) as template"

        # Reusing the read endpoint logic...
        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"):
                record = data["results"][0]
                template_data = record["template"]
                # Deserialize the schema string
                if 'json_schema_str' in template_data:
                    template_data['json_schema'] = json.loads(template_data.pop('json_schema_str'))
                return template_data
            return None
        except Exception as e:
            logger.error(f"Failed to retrieve template: {e}", exc_info=True)
            return None

    # ==========================================================================
    # 3. NEW: ACE Curator Methods (Direct Cypher Writes)
    # ==========================================================================

    async def _execute_write_query(self, query: str, parameters: Dict[str, Any]) -> bool:
        """Helper to send Cypher write queries to 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()
            return response.json().get("status") == "success"
        except Exception as e:
            logger.error(f"GraphDB Write Error: {e}")
            return False

    # services/clients/neo4j_client.py

    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}")
        return await self._execute_write_query(query, params)

    # services/clients/neo4j_client.py (Inside update_bullet_scores)

    async def update_bullet_scores(self, bullet_id: str, feedback_type: Literal['helpful', 'harmful']) -> bool:

        # 1. 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:
            return False

        # 2. Construct the single, syntactically correct Cypher query
        # This must be the exact structure implemented in your GraphDBClient.
        query = f"""
        MATCH (b:ContextBullet {{bullet_id: $bullet_id}}) 
        // This is the critical line: one SET, then score update, then comma, then datetime update
        SET {score_clause}, b.last_feedback_at = datetime() 
        RETURN b
        """

        # 3. Execute the write query
        return await self._execute_write_query(query, {"bullet_id": bullet_id})

    async def get_top_bullets_for_template(self, template_id: str, limit: int = 5) -> List[Dict[str, Any]]:
        """
        Retrieves the highest-scored strategies directly from Neo4j, applying an aging/decay factor.
        """
        # Define an aging factor in Cypher (e.g., divide score by days since creation)
        # This is a conceptual implementation; the actual decay logic depends on your business needs.

        query = """
            MATCH (t:DocumentTemplate {template_id: $template_id})-[:HAS_BULLET]->(b:ContextBullet)
            WITH b,
                 // 🟢 FIX: Ensure helpful_count and harmful_count are coalesced to 0
                 (coalesce(b.helpful_count, 0) - coalesce(b.harmful_count, 0)) AS score_utility,

                 // Using toFloat for compatibility (though duration.days returns a number)
                 toFloat(duration.between(coalesce(b.created_at, datetime()), datetime()).days) AS days_old

            // Ensure 1.0 + ... does not result in division by zero or a problematic float.
            WITH b, score_utility / (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, "limit": limit}
        }

        try:
            response = await self.client.post("/query-cypher", json=payload, timeout=30.0)
            response.raise_for_status()
            results = response.json().get("results", [])

            # 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)
                    # Synthetic score for direct retrieval
                }
                for rec in results
            ]
        except Exception as e:
            logger.error(f"Failed to fetch top bullets: {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,
            query: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """
        Retrieve infrastructure facts for the given template and keywords.

        Args:
            template_id: ID of the document template
            keywords: List of keywords to search for (preferred)
            query: Optional query string (used only if keywords not provided)

        Returns:
            List of infrastructure facts matching the criteria
        """
        if not keywords and query:
            # Fallback to extract keywords if not provided
            logger.warning("Keywords not provided, falling back to query-based extraction")
            # keywords = await self._extract_search_keywords(query)

        if not keywords:
            logger.warning("No keywords provided and no query to extract from")
            return []

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

        # 🟢 OPTIMIZED: Added DISTINCT and toString() for date formatting
        cypher_query = """
        MATCH (t:DocumentTemplate {template_id: $template_id})
        OPTIONAL MATCH (t)-[:HAS_RECENT_FACT]->(f:InfrastructureFact)
        OPTIONAL MATCH (r:Resource)-[:PROVIDES_ASSERTION]->(f)
        WHERE any(kw IN $keywords WHERE 
            toLower(f.content) CONTAINS toLower(kw) OR 
            toLower(r.name) CONTAINS toLower(kw)
        )
        RETURN DISTINCT {
            resource: coalesce(r.name, "unknown_resource"),
            summary: f.content,
            metrics_raw: f.metrics,
            updated: toString(f.updated_at) 
        } as fact
        ORDER BY fact.updated DESC
        LIMIT 5
        """

        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", {})
                # Deserialization remains the same
                metrics_str = fact.get("metrics_raw")
                if metrics_str:
                    try:
                        fact["metrics"] = json.loads(metrics_str)
                    except Exception:
                        fact["metrics"] = {}
                    del fact["metrics_raw"]
                enriched_facts.append(fact)
            return enriched_facts
        except Exception as e:
            logger.error(f"Error fetching facts: {e}")
            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 []

        # ==========================================================================
        # 4. SHARED BLACKBOARD METHODS (Multi-Agent Interop)
        # ==========================================================================

    async def upsert_shared_resources(
            self,
            session_id: str,
            findings: List[Dict[str, Any]],
            template_id: str  # 🟢 Required for contextual anchoring
    ) -> bool:
        if not findings:
            return True

        # 1. Pre-serialize metrics for standard Cypher compatibility
        for finding in findings:
            if isinstance(finding.get("metrics"), dict):
                finding["metrics_json"] = json.dumps(finding["metrics"])
            else:
                finding["metrics_json"] = "{}"

        # 2. Anchored Ingestion Query
        query = """
        UNWIND $findings AS finding

        // Anchor to the Template to ensure domain isolation
        MATCH (t:DocumentTemplate {template_id: $active_template_id})

        // Establish the Resource (The 'What')
        MERGE (r:Resource {name: finding.tool_name})
        SET r.last_observed = datetime(),
            r.current_status = finding.status

        // Establish the Fact (The 'Intelligence') 
        // We unique-ify the Fact ID using the Template ID to prevent cross-talk
        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()

        // Create the semantic links
        MERGE (r)-[:PROVIDES_ASSERTION]->(f)
        MERGE (t)-[:HAS_RECENT_FACT]->(f)

        RETURN count(f) as fact_count
        """

        params = {
            "session_id": session_id,
            "findings": findings,
            "active_template_id": template_id  # 🟢 Mapped to Cypher variable
        }

        logger.info(f"🛰️ Promoting facts for {template_id} to Shared Blackboard.")
        return await self._execute_write_query(query, params)
