#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
Knowledge Graph Query Extensions: chain-finding and hub/stats queries.

Extracted from knowledge_graph_query.py to keep both files under 300 lines.
Contains: query_chain, get_hub_nodes, get_graph_stats (and their SQL/JSON variants).
"""

from collections import defaultdict

from pulse.core.brain_utils import GRAPH_FILE, load_json_dict, search_relevance
from pulse.graph.graph_helpers import (
    _get_conn, _sqlite_has_data, _find_nodes_sql,
    _load_graph_json, _find_matching_nodes_json,
)

import logging
log = logging.getLogger("pulse.graph.query")

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)."""
    conn = _get_conn()
    if _sqlite_has_data(conn):
        return _query_chain_sql(conn, from_text, to_text, max_depth)
    return _query_chain_json(from_text, to_text, max_depth)


def _query_chain_sql(conn, from_text: str, to_text: str, max_depth: int) -> list:
    from_nodes = _find_nodes_sql(conn, from_text, limit=3)
    to_nodes = _find_nodes_sql(conn, to_text, limit=3)
    if not from_nodes or not to_nodes:
        return []
    to_ids = {n["id"] for n in to_nodes}

    for start in from_nodes:
        queue = [(start["id"], [start["id"]])]
        visited = {start["id"]}
        while queue:
            current, path = queue.pop(0)
            if len(path) > max_depth:
                continue
            if current in to_ids:
                chain = []
                for nid in path:
                    row = conn.execute(
                        "SELECT id, type, text FROM graph_nodes WHERE id = ?", (nid,)
                    ).fetchone()
                    chain.append({
                        "node": row["text"] if row else nid,
                        "type": row["type"] if row else "?",
                    })
                return chain
            neighbors = conn.execute(
                "SELECT to_node FROM graph_edges WHERE from_node = ?", (current,)
            ).fetchall()
            for nb in neighbors:
                if nb["to_node"] not in visited:
                    visited.add(nb["to_node"])
                    queue.append((nb["to_node"], path + [nb["to_node"]]))
    return []


def _query_chain_json(from_text: str, to_text: str, max_depth: int) -> list:
    graph = _load_graph_json()
    from_nodes = _find_matching_nodes_json(from_text, graph)
    to_nodes = set(_find_matching_nodes_json(to_text, graph))
    if not from_nodes or not to_nodes:
        return []
    adj = defaultdict(list)
    for edge in graph.get("edges", []):
        adj[edge["from"]].append((edge["to"], edge["type"]))
    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:
                chain = []
                for nid in path:
                    node = graph["nodes"].get(nid, {})
                    chain.append({"node": node.get("text", nid), "type": node.get("type", "?")})
                return chain
            for neighbor, _ 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)."""
    conn = _get_conn()
    if _sqlite_has_data(conn):
        return _get_hub_nodes_sql(conn, top_n)
    return _get_hub_nodes_json(top_n)


def _get_hub_nodes_sql(conn, top_n: int) -> list:
    rows = conn.execute("""
        SELECT n.id, n.type, n.text, COUNT(*) as degree
        FROM graph_nodes n
        JOIN (
            SELECT from_node as node FROM graph_edges
            UNION ALL
            SELECT to_node as node FROM graph_edges
        ) e ON e.node = n.id
        GROUP BY n.id
        ORDER BY degree DESC
        LIMIT ?
    """, (top_n,)).fetchall()
    return [{
        "text": r["text"] or r["id"],
        "type": r["type"] or "?",
        "connections": r["degree"],
        "source": "knowledge_graph",
    } for r in rows]


def _get_hub_nodes_json(top_n: int) -> list:
    graph = _load_graph_json()
    nodes = graph.get("nodes", {})
    edges = graph.get("edges", [])
    degree = defaultdict(int)
    for edge in edges:
        degree[edge["from"]] += 1
        degree[edge["to"]] += 1
    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."""
    conn = _get_conn()
    if _sqlite_has_data(conn):
        return _get_graph_stats_sql(conn)
    return _get_graph_stats_json()


def _get_graph_stats_sql(conn) -> dict:
    node_count = conn.execute("SELECT COUNT(*) as c FROM graph_nodes").fetchone()["c"]
    edge_count = conn.execute("SELECT COUNT(*) as c FROM graph_edges").fetchone()["c"]
    node_types = {}
    for r in conn.execute("SELECT type, COUNT(*) as c FROM graph_nodes GROUP BY type").fetchall():
        node_types[r["type"] or "?"] = r["c"]
    edge_types = {}
    for r in conn.execute("SELECT edge_type, COUNT(*) as c FROM graph_edges GROUP BY edge_type").fetchall():
        edge_types[r["edge_type"] or "?"] = r["c"]
    return {
        "total_nodes": node_count,
        "total_edges": edge_count,
        "node_types": node_types,
        "edge_types": edge_types,
        "built_at": "sqlite",
        "storage": "sqlite",
    }


def _get_graph_stats_json() -> dict:
    graph = _load_graph_json()
    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"),
        "storage": "json_legacy",
    }
