#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
Knowledge Graph Query Engine: Search and path-finding operations.

Provides semantic search over the knowledge graph using causal relationships.

Core functions:
  - query_causes(text)      Find what CAUSES or LEADS_TO the queried concept
  - query_prevents(text)    Find what PREVENTS the queried concept
  - query_related(text)     Find all nodes related to query (any edge type)
  - query_chain(from, to)   Find causal path between two concepts
  - get_hub_nodes(n)        Get most connected knowledge (highest degree)
  - get_graph_stats()       Summary statistics about the graph
"""

import sys
from collections import defaultdict
from pathlib import Path

# Path setup for imports

from pulse.core.brain_utils import KNOWLEDGE_DIR, load_json_dict, search_relevance

GRAPH_FILE = KNOWLEDGE_DIR / "knowledge_graph.json"


def _load_graph() -> dict:
    """Load graph from disk."""
    if not GRAPH_FILE.exists():
        return {"nodes": {}, "edges": []}
    return load_json_dict(GRAPH_FILE)


def _find_matching_nodes(query: str, graph: dict, threshold: float = 0.25) -> list:
    """Find graph nodes matching a query by text similarity.

    Returns list of node IDs sorted by relevance.
    """
    matches = []
    nodes = graph.get("nodes", {})
    for nid, node in nodes.items():
        text = node.get("text", "")
        sim = search_relevance(query, text)
        if sim >= threshold:
            matches.append((nid, sim))

    matches.sort(key=lambda x: -x[1])
    return [nid for nid, _ in matches[:10]]


def query_causes(query: str, max_results: int = 5) -> list:
    """Find what CAUSES or LEADS_TO the queried concept.

    Returns list of dicts with text, type, relationship, confidence, source.
    """
    graph = _load_graph()
    results = []

    # Find nodes matching the query
    matching_nodes = _find_matching_nodes(query, graph)

    for node_id in matching_nodes:
        # Find incoming edges (things that cause this node)
        for edge in graph.get("edges", []):
            if edge["to"] == node_id and edge["type"] in ("CAUSES", "LEADS_TO"):
                from_node = graph["nodes"].get(edge["from"], {})
                if from_node:
                    results.append({
                        "text": from_node.get("text", ""),
                        "type": from_node.get("type", "CONCEPT"),
                        "relationship": edge["type"],
                        "confidence": edge.get("confidence", 0.5),
                        "source": "knowledge_graph",
                    })

    return results[:max_results]


def query_prevents(query: str, max_results: int = 5) -> list:
    """Find what PREVENTS the queried concept.

    Returns list of dicts with text, type, relationship, confidence, source.
    """
    graph = _load_graph()
    results = []

    matching_nodes = _find_matching_nodes(query, graph)

    for node_id in matching_nodes:
        for edge in graph.get("edges", []):
            if edge["to"] == node_id and edge["type"] == "PREVENTS":
                from_node = graph["nodes"].get(edge["from"], {})
                if from_node:
                    results.append({
                        "text": from_node.get("text", ""),
                        "type": from_node.get("type", "CONCEPT"),
                        "relationship": "PREVENTS",
                        "confidence": edge.get("confidence", 0.5),
                        "source": "knowledge_graph",
                    })

    return results[:max_results]


def query_related(query: str, max_results: int = 5) -> list:
    """General query: find all nodes related to the query (any edge type).

    Returns list of dicts with text, type, relationship, confidence, source.
    """
    graph = _load_graph()
    results = []

    matching_nodes = _find_matching_nodes(query, graph)

    # Collect all connected nodes (1-hop neighbors)
    neighbor_ids = set()
    for node_id in matching_nodes:
        for edge in graph.get("edges", []):
            if edge["from"] == node_id:
                neighbor_ids.add((edge["to"], edge["type"], edge.get("confidence", 0.5)))
            elif edge["to"] == node_id:
                neighbor_ids.add((edge["from"], edge["type"], edge.get("confidence", 0.5)))

    for nid, rel_type, conf in neighbor_ids:
        node = graph["nodes"].get(nid, {})
        if node:
            results.append({
                "text": node.get("text", ""),
                "type": node.get("type", "CONCEPT"),
                "relationship": rel_type,
                "confidence": conf,
                "source": "knowledge_graph",
            })

    # Sort by confidence
    results.sort(key=lambda x: -x["confidence"])
    return results[:max_results]


def query_chain(from_text: str, to_text: str, max_depth: int = 4) -> list:
    """Find causal path between two concepts (BFS up to max_depth).

    Returns list of dicts with node, type for each step in the causal chain.
    Empty list if no path found.
    """
    graph = _load_graph()

    from_nodes = _find_matching_nodes(from_text, graph)
    to_nodes = set(_find_matching_nodes(to_text, graph))

    if not from_nodes or not to_nodes:
        return []

    # Build adjacency list
    adj = defaultdict(list)
    for edge in graph.get("edges", []):
        adj[edge["from"]].append((edge["to"], edge["type"]))

    # BFS from start nodes
    for start in from_nodes:
        queue = [(start, [start])]
        visited = {start}

        while queue:
            current, path = queue.pop(0)
            if len(path) > max_depth:
                continue

            if current in to_nodes:
                # Build readable chain
                chain = []
                for i, nid in enumerate(path):
                    node = graph["nodes"].get(nid, {})
                    chain.append({
                        "node": node.get("text", nid),
                        "type": node.get("type", "?"),
                    })
                return chain

            for neighbor, rel_type in adj.get(current, []):
                if neighbor not in visited:
                    visited.add(neighbor)
                    queue.append((neighbor, path + [neighbor]))

    return []


def get_hub_nodes(top_n: int = 10) -> list:
    """Get the most connected nodes (highest degree centrality).

    Returns list of dicts with text, type, connections, source.
    """
    graph = _load_graph()
    nodes = graph.get("nodes", {})
    edges = graph.get("edges", [])

    # Count connections per node
    degree = defaultdict(int)
    for edge in edges:
        degree[edge["from"]] += 1
        degree[edge["to"]] += 1

    # Sort by degree
    hubs = sorted(degree.items(), key=lambda x: -x[1])

    results = []
    for nid, deg in hubs[:top_n]:
        node = nodes.get(nid, {})
        results.append({
            "text": node.get("text", nid),
            "type": node.get("type", "?"),
            "connections": deg,
            "source": "knowledge_graph",
        })

    return results


def get_graph_stats() -> dict:
    """Get summary statistics about the knowledge graph.

    Returns dict with total_nodes, total_edges, node_types, edge_types, built_at.
    """
    graph = _load_graph()
    nodes = graph.get("nodes", {})
    edges = graph.get("edges", [])

    type_counts = defaultdict(int)
    for n in nodes.values():
        type_counts[n.get("type", "?")] += 1

    edge_type_counts = defaultdict(int)
    for e in edges:
        edge_type_counts[e.get("type", "?")] += 1

    return {
        "total_nodes": len(nodes),
        "total_edges": len(edges),
        "node_types": dict(type_counts),
        "edge_types": dict(edge_type_counts),
        "built_at": graph.get("metadata", {}).get("built_at", "never"),
    }
