#!/usr/bin/env python3
"""Unit tests: knowledge graph node degree computation after rebuild.

Audits two degree calculation paths:
  1. Stored column: UPDATE graph_nodes SET degree = ... (sync_to_sqlite)
  2. Query-time JOIN: COUNT(*) via UNION ALL (get_hub_nodes)

All tests use isolated in-memory SQLite — no production DB writes.
Avoids the anti-pattern: dangling edges from tests hitting real DB.
"""
import sqlite3
import unittest


# ---------------------------------------------------------------------------
# In-memory DB helpers
# ---------------------------------------------------------------------------

def _make_in_memory_db() -> sqlite3.Connection:
    """Create an isolated in-memory SQLite DB with the graph schema from brain_db.py."""
    conn = sqlite3.connect(":memory:")
    conn.row_factory = sqlite3.Row
    conn.execute("""
        CREATE TABLE graph_nodes (
            id TEXT PRIMARY KEY, type TEXT NOT NULL DEFAULT 'CONCEPT',
            text TEXT NOT NULL, source TEXT DEFAULT '', confidence REAL DEFAULT 0.5,
            domain TEXT NOT NULL DEFAULT 'general',
            date TEXT DEFAULT '', degree INTEGER DEFAULT 0, embedding BLOB DEFAULT NULL,
            metadata TEXT DEFAULT '{}'
        )
    """)
    conn.execute("""
        CREATE TABLE graph_edges (
            id INTEGER PRIMARY KEY AUTOINCREMENT, from_node TEXT NOT NULL,
            to_node TEXT NOT NULL, edge_type TEXT NOT NULL,
            confidence REAL DEFAULT 0.5, source_node TEXT DEFAULT '',
            metadata TEXT DEFAULT '{}',
            UNIQUE(from_node, edge_type, to_node)
        )
    """)
    conn.commit()
    return conn


def _insert_node(conn: sqlite3.Connection, nid: str, ntype: str = "CONCEPT", text: str = ""):
    conn.execute(
        "INSERT INTO graph_nodes (id, type, text) VALUES (?, ?, ?)",
        (nid, ntype, text or nid),
    )


def _insert_edge(conn: sqlite3.Connection, from_node: str, to_node: str, edge_type: str = "LEADS_TO"):
    conn.execute(
        "INSERT INTO graph_edges (from_node, to_node, edge_type) VALUES (?, ?, ?)",
        (from_node, to_node, edge_type),
    )


def _run_degree_update(conn: sqlite3.Connection):
    """Run the exact UPDATE used in graph_sqlite_sync.sync_to_sqlite."""
    conn.execute("""
        UPDATE graph_nodes SET degree = (
            SELECT COUNT(*) FROM graph_edges
            WHERE from_node = graph_nodes.id
               OR to_node = graph_nodes.id
        )
    """)
    conn.commit()


def _get_hub_nodes_sql(conn: sqlite3.Connection, top_n: int = 10) -> list:
    """Mirror of knowledge_graph_query._get_hub_nodes_sql."""
    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 [dict(r) for r in rows]


# ---------------------------------------------------------------------------
# Tests: stored degree column (sync_to_sqlite UPDATE)
# ---------------------------------------------------------------------------

class TestDegreeColumnAfterUpdate(unittest.TestCase):
    """Stored degree column correctness after rebuild UPDATE."""

    def setUp(self):
        self.conn = _make_in_memory_db()

    def tearDown(self):
        self.conn.close()

    def test_isolated_nodes_default_to_zero_degree(self):
        _insert_node(self.conn, "A")
        _insert_node(self.conn, "B")
        self.conn.commit()
        rows = self.conn.execute("SELECT degree FROM graph_nodes").fetchall()
        self.assertTrue(all(r["degree"] == 0 for r in rows),
                        "Freshly inserted nodes must default to degree=0")

    def test_degree_updated_correctly_for_hub_node(self):
        """Hub with 3 edges must have degree=3 after UPDATE."""
        for nid in ("hub", "leaf1", "leaf2", "leaf3"):
            _insert_node(self.conn, nid)
        self.conn.commit()
        _insert_edge(self.conn, "hub", "leaf1")
        _insert_edge(self.conn, "hub", "leaf2")
        _insert_edge(self.conn, "leaf3", "hub")
        self.conn.commit()

        _run_degree_update(self.conn)

        hub_degree = self.conn.execute(
            "SELECT degree FROM graph_nodes WHERE id = 'hub'"
        ).fetchone()["degree"]
        self.assertEqual(hub_degree, 3, "Hub connected to 3 edges should have degree=3")

    def test_leaf_nodes_have_degree_one(self):
        _insert_node(self.conn, "A")
        _insert_node(self.conn, "B")
        self.conn.commit()
        _insert_edge(self.conn, "A", "B")
        self.conn.commit()
        _run_degree_update(self.conn)

        for nid in ("A", "B"):
            deg = self.conn.execute(
                "SELECT degree FROM graph_nodes WHERE id = ?", (nid,)
            ).fetchone()["degree"]
            self.assertEqual(deg, 1, f"Node {nid} connected by one edge should have degree=1")

    def test_degree_counts_both_from_and_to_directions(self):
        """Degree is undirected: counts edges where node is from_node OR to_node."""
        for nid in ("center", "in_node", "out_node"):
            _insert_node(self.conn, nid)
        self.conn.commit()
        _insert_edge(self.conn, "in_node", "center")   # center is to_node
        _insert_edge(self.conn, "center", "out_node")  # center is from_node
        self.conn.commit()
        _run_degree_update(self.conn)

        center_degree = self.conn.execute(
            "SELECT degree FROM graph_nodes WHERE id = 'center'"
        ).fetchone()["degree"]
        self.assertEqual(center_degree, 2, "Center has 1 inbound + 1 outbound = degree 2")

    def test_isolated_nodes_remain_zero_after_update(self):
        _insert_node(self.conn, "orphan")
        _insert_node(self.conn, "connected_a")
        _insert_node(self.conn, "connected_b")
        self.conn.commit()
        _insert_edge(self.conn, "connected_a", "connected_b")
        self.conn.commit()
        _run_degree_update(self.conn)

        orphan_degree = self.conn.execute(
            "SELECT degree FROM graph_nodes WHERE id = 'orphan'"
        ).fetchone()["degree"]
        self.assertEqual(orphan_degree, 0, "Orphan node must retain degree=0 after UPDATE")

    def test_regression_no_connected_node_has_zero_degree_after_rebuild(self):
        """Regression: degree column must NOT stay at 0 after rebuild (was bug pre-fix)."""
        for i in range(5):
            _insert_node(self.conn, f"n{i}")
        self.conn.commit()
        for i in range(4):
            _insert_edge(self.conn, f"n{i}", f"n{i + 1}")
        self.conn.commit()
        _run_degree_update(self.conn)

        zero_degree_connected = self.conn.execute("""
            SELECT COUNT(*) FROM graph_nodes
            WHERE degree = 0
              AND id IN (
                  SELECT from_node FROM graph_edges
                  UNION SELECT to_node FROM graph_edges
              )
        """).fetchone()[0]
        self.assertEqual(
            zero_degree_connected, 0,
            f"{zero_degree_connected} connected node(s) still have degree=0 after UPDATE — "
            "regression: missing UPDATE in sync_to_sqlite",
        )

    def test_degree_update_is_idempotent(self):
        """Running degree UPDATE twice must produce the same result."""
        for nid in ("X", "Y", "Z"):
            _insert_node(self.conn, nid)
        self.conn.commit()
        _insert_edge(self.conn, "X", "Y")
        _insert_edge(self.conn, "Y", "Z")
        self.conn.commit()

        _run_degree_update(self.conn)
        first = {
            r["id"]: r["degree"]
            for r in self.conn.execute("SELECT id, degree FROM graph_nodes").fetchall()
        }
        _run_degree_update(self.conn)
        second = {
            r["id"]: r["degree"]
            for r in self.conn.execute("SELECT id, degree FROM graph_nodes").fetchall()
        }
        self.assertEqual(first, second, "Degree UPDATE must be idempotent")


# ---------------------------------------------------------------------------
# Tests: query-time degree via UNION ALL JOIN (get_hub_nodes)
# ---------------------------------------------------------------------------

class TestQueryTimeDegreeJoin(unittest.TestCase):
    """Query-time degree correctness via UNION ALL JOIN in _get_hub_nodes_sql."""

    def setUp(self):
        self.conn = _make_in_memory_db()

    def tearDown(self):
        self.conn.close()

    def test_hub_nodes_ordered_by_degree_desc(self):
        """get_hub_nodes must return nodes in descending degree order."""
        for nid in ("hub", "mid", "leaf"):
            _insert_node(self.conn, nid)
        self.conn.commit()
        # hub: 2 connections, mid: 2, leaf: 2 (but hub highest after UNION ALL)
        # Make it unambiguous: hub→leaf, hub→mid, mid→leaf
        _insert_edge(self.conn, "hub", "leaf")
        _insert_edge(self.conn, "hub", "mid")
        _insert_edge(self.conn, "mid", "leaf")
        self.conn.commit()

        hubs = _get_hub_nodes_sql(self.conn, top_n=3)
        self.assertEqual(len(hubs), 3)
        degrees = [h["degree"] for h in hubs]
        self.assertEqual(
            degrees, sorted(degrees, reverse=True),
            f"Hub nodes not in descending degree order: {degrees}",
        )

    def test_query_time_degree_matches_stored_degree(self):
        """Stored degree column and query-time JOIN degree must agree after UPDATE."""
        for nid in ("A", "B", "C", "D"):
            _insert_node(self.conn, nid)
        self.conn.commit()
        _insert_edge(self.conn, "A", "B")
        _insert_edge(self.conn, "A", "C")
        _insert_edge(self.conn, "A", "D")
        _insert_edge(self.conn, "B", "C")
        self.conn.commit()
        _run_degree_update(self.conn)

        hubs = _get_hub_nodes_sql(self.conn, top_n=10)
        hub_map = {h["id"]: h["degree"] for h in hubs}

        for nid in ("A", "B", "C", "D"):
            stored = self.conn.execute(
                "SELECT degree FROM graph_nodes WHERE id = ?", (nid,)
            ).fetchone()["degree"]
            query_time = hub_map.get(nid, 0)
            self.assertEqual(
                stored, query_time,
                f"Node {nid}: stored degree={stored} != query-time degree={query_time}",
            )

    def test_top_n_limit_respected(self):
        """get_hub_nodes respects the top_n limit."""
        for i in range(20):
            _insert_node(self.conn, f"n{i}")
        self.conn.commit()
        for i in range(1, 20):
            _insert_edge(self.conn, "n0", f"n{i}")
        self.conn.commit()

        hubs = _get_hub_nodes_sql(self.conn, top_n=5)
        self.assertLessEqual(len(hubs), 5, "get_hub_nodes must honor top_n limit")

    def test_disconnected_nodes_absent_from_hub_results(self):
        """Nodes with no edges must NOT appear in hub node results."""
        _insert_node(self.conn, "isolated")
        _insert_node(self.conn, "connected_a")
        _insert_node(self.conn, "connected_b")
        self.conn.commit()
        _insert_edge(self.conn, "connected_a", "connected_b")
        self.conn.commit()

        hubs = _get_hub_nodes_sql(self.conn, top_n=10)
        hub_ids = {h["id"] for h in hubs}
        self.assertNotIn("isolated", hub_ids,
                         "Isolated nodes must not appear in hub node results")

    def test_multiple_edge_types_counted_correctly(self):
        """All edge types count toward degree (no filter by edge_type)."""
        for nid in ("center", "a", "b", "c"):
            _insert_node(self.conn, nid)
        self.conn.commit()
        _insert_edge(self.conn, "center", "a", "LEADS_TO")
        _insert_edge(self.conn, "center", "b", "PREVENTS")
        _insert_edge(self.conn, "center", "c", "REQUIRES")
        self.conn.commit()
        _run_degree_update(self.conn)

        stored = self.conn.execute(
            "SELECT degree FROM graph_nodes WHERE id = 'center'"
        ).fetchone()["degree"]
        self.assertEqual(stored, 3, "All 3 edge types must count toward degree")

        hubs = _get_hub_nodes_sql(self.conn, top_n=1)
        self.assertEqual(hubs[0]["id"], "center")
        self.assertEqual(hubs[0]["degree"], 3)

    def test_self_loop_behavior_documents_method_difference(self):
        """Documents: stored UPDATE counts edge rows (1), UNION ALL counts appearances (2).

        Self-loops don't occur in the knowledge graph builder, so this is academic.
        The stored degree UPDATE uses WHERE ... OR ..., counting the row once.
        The query-time UNION ALL counts the node appearing in both from and to.
        Both methods agree for all real-world cases (no self-loops).
        """
        _insert_node(self.conn, "self")
        self.conn.commit()
        _insert_edge(self.conn, "self", "self")  # self-loop (not produced by builder)
        self.conn.commit()
        _run_degree_update(self.conn)

        stored = self.conn.execute(
            "SELECT degree FROM graph_nodes WHERE id = 'self'"
        ).fetchone()["degree"]
        # Stored UPDATE: WHERE from_node='self' OR to_node='self' → COUNT(*)=1 (one row)
        self.assertEqual(stored, 1,
                         "Stored UPDATE counts 1 row for self-loop (OR condition, one row)")

        hubs = _get_hub_nodes_sql(self.conn, top_n=1)
        query_time = hubs[0]["degree"] if hubs else 0
        # UNION ALL: 'self' appears in from_node + to_node → counted twice
        self.assertEqual(query_time, 2,
                         "UNION ALL counts self-loop node twice (directional appearances)")


if __name__ == "__main__":
    unittest.main(verbosity=2)
