# Copyright (c) 2026 Varun Pratap Bhardwaj / Qualixar
# Licensed under AGPL-3.0-or-later - see LICENSE file
# Part of SuperLocalMemory V3 | https://qualixar.com | https://varunpratap.com

"""SuperLocalMemory V3 — Entity Graph Channel with Spreading Activation.

SA-RAG pattern: entities from query -> canonical lookup -> graph traversal
with decay. Handles BOTH uppercase and lowercase entity mentions.

Part of Qualixar | Author: Varun Pratap Bhardwaj
License: Elastic-2.0
"""
from __future__ import annotations

import json
import logging
import re
from collections import defaultdict
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from superlocalmemory.encoding.entity_resolver import EntityResolver
    from superlocalmemory.storage.database import DatabaseManager

logger = logging.getLogger(__name__)

_PROPER_NOUN_RE = re.compile(r"\b[A-Z][a-z]{1,}\b")

_ENTITY_STOP: frozenset[str] = frozenset({
    # Expanded stop list for query entity extraction
    "what", "when", "where", "who", "which", "how", "does", "did",
    "the", "that", "this", "there", "then", "than", "they", "them",
    "have", "has", "had", "been", "being", "about", "after", "before",
    "from", "into", "with", "some", "other", "would", "could", "should",
    "will", "because", "also", "just", "like", "know", "think",
    "feel", "want", "need", "make", "take", "give", "tell", "said",
    "wow", "gonna", "got", "by", "thanks", "thank", "hey", "hi",
    "hello", "bye", "good", "great", "nice", "cool", "right",
    "let", "can", "might", "much", "many", "more", "most",
    "something", "anything", "everything", "nothing", "someone",
    "it", "my", "your", "our", "their", "me", "you", "we", "us",
    "do", "if", "or", "no", "to", "at", "on", "in", "so",
    "go", "come", "see", "look", "say", "ask", "try", "keep",
    "yes", "yeah", "sure", "okay", "ok", "really", "actually",
    "maybe", "well", "still", "even", "very",
})


def extract_query_entities(query: str) -> list[str]:
    """Extract entity candidates from query (handles both cases).

    Strategy: find proper nouns in original + title-cased text,
    plus quoted phrases. Deduplicates case-insensitively.
    """
    candidates: list[str] = []
    seen: set[str] = set()

    def _add(name: str) -> None:
        lo = name.lower()
        if lo not in seen and lo not in _ENTITY_STOP and len(name) >= 2:
            seen.add(lo)
            candidates.append(name)

    for m in _PROPER_NOUN_RE.finditer(query):
        _add(m.group(0))
    for m in _PROPER_NOUN_RE.finditer(query.title()):
        _add(m.group(0))
    for m in re.finditer(r'"([^"]+)"', query):
        _add(m.group(1).strip())
    # Also extract multi-word capitalized sequences (e.g. "New York", "San Francisco")
    for m in re.finditer(r'\b([A-Z][a-z]+(?:\s+[A-Z][a-z]+)+)\b', query):
        _add(m.group(1))
    # Extract all-caps abbreviations (e.g. NYU, MIT, UCLA) — min 2 chars
    for m in re.finditer(r'\b([A-Z]{2,})\b', query):
        _add(m.group(1))

    return candidates


class EntityGraphChannel:
    """Entity-based retrieval with spreading activation (SA-RAG).

    V3.3.9: In-memory adjacency cache for O(1) edge lookup.
    Replaces per-node SQLite queries (23ms each) with dict lookup (<0.001ms).
    The cache is loaded once per profile and invalidated on store/edge changes.
    Memory cost: ~18 MB for 232K edges. Zero quality change — same algorithm.
    """

    def __init__(
        self, db: DatabaseManager,
        entity_resolver: EntityResolver | None = None,
        decay: float = 0.7, activation_threshold: float = 0.05,
        max_hops: int = 4,
        graph_metrics: dict[str, dict] | None = None,
        cozo_backend: Any = None,  # v3.4.5: optional CozoDB backend
    ) -> None:
        self._db = db
        self._resolver = entity_resolver
        self._decay = decay
        self._threshold = activation_threshold
        self._max_hops = max_hops
        # v3.4.5: Optional CozoDB graph backend (Sprint 2)
        self._cozo = cozo_backend
        # In-memory adjacency: {node_id -> [(neighbor_id, weight), ...]}
        self._adj: dict[str, list[tuple[str, float]]] = {}
        self._adj_profile: str = ""  # Track which profile is loaded
        self._adj_edge_count: int = 0  # Track edge count for staleness detection
        # v3.4.1: Graph intelligence metrics (loaded from fact_importance)
        self._graph_metrics: dict[str, dict] = graph_metrics or {}
        self._graph_metrics_profile: str = ""

    def _ensure_adjacency(self, profile_id: str) -> None:
        """Load graph adjacency into memory for fast spreading activation.

        Loads ALL edges for a profile into a bidirectional dict.
        Called once per profile switch or when edge count changes (new store).
        Cost: ~1s for 232K edges, ~18 MB RAM.
        """
        # Check staleness: profile changed or new edges added since last load
        current_count = self._get_edge_count(profile_id)
        if (self._adj_profile == profile_id
                and self._adj
                and self._adj_edge_count == current_count):
            return
        adj: dict[str, list[tuple[str, float]]] = defaultdict(list)
        try:
            rows = self._db.execute(
                "SELECT source_id, target_id, weight FROM graph_edges WHERE profile_id = ?",
                (profile_id,),
            )
        except Exception:
            rows = []
        for r in rows:
            d = dict(r)
            s, t, w = d["source_id"], d["target_id"], float(d["weight"])
            adj[s].append((t, w))
            adj[t].append((s, w))
        self._adj = dict(adj)  # Convert defaultdict to regular dict (no accidental growth)
        self._adj_profile = profile_id
        self._adj_edge_count = current_count
        # Also load entity maps (same staleness lifecycle)
        self._load_entity_maps(profile_id)
        # v3.4.1: Load graph intelligence metrics (P0)
        self._load_graph_metrics(profile_id)

        logger.info(
            "Loaded adjacency cache: %d nodes, %d edges, %d entity mappings for profile %s",
            len(self._adj), sum(len(v) for v in self._adj.values()) // 2,
            len(self._entity_to_facts), profile_id,
        )

    def _get_edge_count(self, profile_id: str) -> int:
        """Fast edge count for staleness check (~1ms)."""
        try:
            rows = self._db.execute(
                "SELECT COUNT(*) as cnt FROM graph_edges WHERE profile_id = ?",
                (profile_id,),
            )
            if rows:
                return int(dict(rows[0]).get("cnt", 0))
        except Exception:
            pass
        return 0

    def _load_entity_maps(self, profile_id: str) -> None:
        """Pre-load entity→fact and fact→entity maps into memory.

        Eliminates per-entity and per-fact SQL in the spreading activation loop.
        Same data, same algorithm — zero quality change.
        """
        # entity_id -> [fact_id, ...]
        self._entity_to_facts: dict[str, list[str]] = defaultdict(list)
        # fact_id -> [entity_id, ...]
        self._fact_to_entities: dict[str, list[str]] = defaultdict(list)

        try:
            rows = self._db.execute(
                "SELECT fact_id, canonical_entities_json FROM atomic_facts "
                "WHERE profile_id = ? AND canonical_entities_json IS NOT NULL "
                "AND canonical_entities_json != ''",
                (profile_id,),
            )
        except Exception:
            rows = []
        for r in rows:
            d = dict(r)
            fid = d["fact_id"]
            raw = d.get("canonical_entities_json")
            if not raw:
                continue
            try:
                eids = json.loads(raw)
                for eid in eids:
                    self._entity_to_facts[eid].append(fid)
                    self._fact_to_entities[fid].append(eid)
            except (ValueError, TypeError):
                continue

        logger.info(
            "Loaded entity maps: %d entities, %d facts with entities",
            len(self._entity_to_facts), len(self._fact_to_entities),
        )

    def _load_graph_metrics(self, profile_id: str) -> None:
        """Load PageRank, community_id, degree_centrality from fact_importance.

        v3.4.1: Enables graph-enhanced retrieval (P0).
        Called alongside adjacency loading. Same staleness lifecycle.
        """
        if self._graph_metrics_profile == profile_id and self._graph_metrics:
            return
        self._graph_metrics = {}
        self._graph_metrics_profile = profile_id
        try:
            rows = self._db.execute(
                "SELECT fact_id, pagerank_score, community_id, degree_centrality "
                "FROM fact_importance WHERE profile_id = ?",
                (profile_id,),
            )
            for r in rows:
                d = dict(r)
                self._graph_metrics[d["fact_id"]] = {
                    "pagerank_score": float(d.get("pagerank_score", 0) or 0),
                    "community_id": d.get("community_id"),
                    "degree_centrality": float(d.get("degree_centrality", 0) or 0),
                }
            logger.info(
                "Loaded graph metrics: %d facts for profile %s",
                len(self._graph_metrics), profile_id,
            )
        except Exception as exc:
            logger.debug("Graph metrics load failed (graceful degradation): %s", exc)
            self._graph_metrics = {}

    def invalidate_cache(self) -> None:
        """Clear all caches. Call after adding/removing edges or facts."""
        self._adj.clear()
        self._adj_profile = ""
        self._adj_edge_count = 0
        self._entity_to_facts = defaultdict(list)
        self._fact_to_entities = defaultdict(list)
        self._graph_metrics.clear()
        self._graph_metrics_profile = ""

    def search(self, query: str, profile_id: str, top_k: int = 50) -> list[tuple[str, float]]:
        """Search via entity graph with spreading activation.

        V3.3.9: Uses in-memory adjacency for O(1) edge lookups.
        V3.4.5: Routes to CozoDB if backend is active (Sprint 2).
        """
        raw_entities = extract_query_entities(query)

        # v3.4.5: Route to CozoDB if active
        if self._cozo is not None:
            return self._search_via_cozo(query, raw_entities, profile_id, top_k)
        if not raw_entities:
            return []

        canonical_ids = self._resolve_entities(raw_entities, profile_id)
        if not canonical_ids:
            return []

        # Load adjacency cache (no-op if already loaded for this profile)
        self._ensure_adjacency(profile_id)

        # Seed activation from direct entity-linked facts
        # Use in-memory map when available, fall back to SQL for mock/test DBs
        activation: dict[str, float] = defaultdict(float)
        visited_entities: set[str] = set(canonical_ids)

        use_cache = bool(self._entity_to_facts)
        for eid in canonical_ids:
            if use_cache:
                for fid in self._entity_to_facts.get(eid, ()):
                    activation[fid] = max(activation[fid], 1.0)
            else:
                for fact in self._db.get_facts_by_entity(eid, profile_id):
                    activation[fact.fact_id] = max(activation[fact.fact_id], 1.0)

        # Spreading activation through graph edges (all in-memory O(1) lookups)
        frontier = set(activation.keys())
        for hop in range(1, self._max_hops):
            hop_decay = self._decay ** hop
            if hop_decay < self._threshold:
                break
            next_frontier: set[str] = set()

            for fid in frontier:
                if use_cache:
                    neighbors = self._adj.get(fid, ())
                    for neighbor, edge_weight in neighbors:
                        # v3.4.2: Only apply edge_weight and PageRank bias when
                        # graph metrics are available. Without metrics, edge_weight
                        # dampens propagation by ~14% with no compensating boost,
                        # causing retrieval regression (68.4% vs 70.4% on LoCoMo).
                        if self._graph_metrics:
                            weighted = activation[fid] * self._decay * edge_weight
                            if neighbor in self._graph_metrics:
                                target_pr = self._graph_metrics[neighbor].get("pagerank_score", 0.0)
                                pr_boost = min(1.0 + target_pr * 2.0, 2.0)
                                weighted *= pr_boost
                        else:
                            weighted = activation[fid] * self._decay
                        if weighted >= self._threshold and weighted > activation.get(neighbor, 0.0):
                            activation[neighbor] = weighted
                            next_frontier.add(neighbor)
                else:
                    # NOTE: SQL fallback path does NOT use graph intelligence (P1/P2/P3).
                    # Graph intelligence is only available on the in-memory cache path.
                    # This fallback exists for mock/test DBs. See Phase 7 LLD H-01.
                    for edge in self._db.get_edges_for_node(fid, profile_id):
                        neighbor = edge.target_id if edge.source_id == fid else edge.source_id
                        propagated = activation[fid] * self._decay
                        if propagated >= self._threshold and propagated > activation.get(neighbor, 0.0):
                            activation[neighbor] = propagated
                            next_frontier.add(neighbor)

            # Discover new entities from activated facts
            if use_cache:
                new_eids: list[str] = []
                for fid in frontier:
                    for eid in self._fact_to_entities.get(fid, ()):
                        if eid not in visited_entities:
                            visited_entities.add(eid)
                            new_eids.append(eid)
                for eid in new_eids:
                    for fid in self._entity_to_facts.get(eid, ()):
                        if hop_decay > activation.get(fid, 0.0):
                            activation[fid] = hop_decay
                            next_frontier.add(fid)
            else:
                # SQL fallback (mock/test DBs)
                new_eids_sql = self._discover_entities(frontier, profile_id, visited_entities)
                for eid in new_eids_sql:
                    visited_entities.add(eid)
                    for fact in self._db.get_facts_by_entity(eid, profile_id):
                        if hop_decay > activation.get(fact.fact_id, 0.0):
                            activation[fact.fact_id] = hop_decay
                            next_frontier.add(fact.fact_id)

            frontier = next_frontier
            if not frontier:
                break

        # v3.4.1 P2: Community-aware boosting
        if self._graph_metrics and use_cache:
            from collections import Counter as _Counter
            seed_communities: _Counter = _Counter()
            for eid in canonical_ids:
                for fid in self._entity_to_facts.get(eid, ()):
                    m = self._graph_metrics.get(fid, {})
                    comm = m.get("community_id")
                    if comm is not None:
                        seed_communities[comm] += 1
            if seed_communities:
                total_seeds = sum(seed_communities.values())
                for fid in list(activation.keys()):
                    m = self._graph_metrics.get(fid, {})
                    fact_comm = m.get("community_id")
                    if fact_comm is not None and fact_comm in seed_communities:
                        boost = min(1.0 + 0.15 * (seed_communities[fact_comm] / total_seeds), 1.3)
                        activation[fid] *= boost
                    elif fact_comm is not None and fact_comm not in seed_communities:
                        activation[fid] *= 0.9  # Mild penalty for unrelated communities

        # v3.4.1 P3: Contradiction suppression via graph_edges
        if use_cache and activation:
            self._suppress_contradictions(activation, profile_id)

        # v3.4.1: Score normalization to [0, 1]
        results = [(fid, sc) for fid, sc in activation.items() if sc >= self._threshold]
        if not results:
            return []
        max_score = max(sc for _, sc in results)
        if max_score > 0:
            results = [(fid, sc / max_score) for fid, sc in results]
        results.sort(key=lambda x: x[1], reverse=True)
        return results[:top_k]

    def score_candidates(
        self,
        query: str,
        candidate_fact_ids: list[str],
        profile_id: str,
    ) -> dict[str, float]:
        """Score candidate facts by their entity-graph proximity to query entities.

        V3.4.11 "Signal Enhancer" architecture: instead of returning its own
        independent set of fact_ids (which get outranked by multi-channel facts
        in RRF), this method scores EXISTING candidates from semantic/BM25
        by their graph connectivity to query entities.

        Research basis: Microsoft GraphRAG DRIFT Search, HippoRAG, Pistis-RAG
        cascaded architecture. Graph signals act as post-retrieval boosters,
        not independent retrievers. Avoids the "weakest link" phenomenon where
        non-overlapping result sets cause rank collapse in RRF fusion.

        Args:
            query: The user's query string.
            candidate_fact_ids: Fact IDs from semantic/BM25/other channels.
            profile_id: User profile.

        Returns:
            Dict mapping fact_id → entity_graph score [0, 1].
            Facts with no entity connection return 0.
            Facts directly linked to query entities score ~1.0.
            Facts 1-hop away score ~0.7 (decay factor).
        """
        if not candidate_fact_ids:
            return {}

        raw_entities = extract_query_entities(query)
        if not raw_entities:
            return {}

        canonical_ids = self._resolve_entities(raw_entities, profile_id)
        if not canonical_ids:
            return {}

        self._ensure_adjacency(profile_id)

        # Run full spreading activation (same as search())
        activation: dict[str, float] = defaultdict(float)
        visited_entities: set[str] = set(canonical_ids)
        use_cache = bool(self._entity_to_facts)

        for eid in canonical_ids:
            if use_cache:
                for fid in self._entity_to_facts.get(eid, ()):
                    activation[fid] = max(activation[fid], 1.0)
            else:
                for fact in self._db.get_facts_by_entity(eid, profile_id):
                    activation[fact.fact_id] = max(activation[fact.fact_id], 1.0)

        frontier = set(activation.keys())
        for hop in range(1, self._max_hops):
            hop_decay = self._decay ** hop
            if hop_decay < self._threshold:
                break
            next_frontier: set[str] = set()
            for fid in frontier:
                if use_cache:
                    for neighbor, edge_weight in self._adj.get(fid, ()):
                        if self._graph_metrics:
                            weighted = activation[fid] * self._decay * edge_weight
                            if neighbor in self._graph_metrics:
                                pr = self._graph_metrics[neighbor].get("pagerank_score", 0.0)
                                weighted *= min(1.0 + pr * 2.0, 2.0)
                        else:
                            weighted = activation[fid] * self._decay
                        if weighted >= self._threshold and weighted > activation.get(neighbor, 0.0):
                            activation[neighbor] = weighted
                            next_frontier.add(neighbor)

            if use_cache:
                for fid in frontier:
                    for eid in self._fact_to_entities.get(fid, ()):
                        if eid not in visited_entities:
                            visited_entities.add(eid)
                            for linked_fid in self._entity_to_facts.get(eid, ()):
                                if hop_decay > activation.get(linked_fid, 0.0):
                                    activation[linked_fid] = hop_decay
                                    next_frontier.add(linked_fid)

            frontier = next_frontier
            if not frontier:
                break

        # Community-aware boosting (same as search)
        if self._graph_metrics and use_cache:
            from collections import Counter as _Counter
            seed_communities: _Counter = _Counter()
            for eid in canonical_ids:
                for fid in self._entity_to_facts.get(eid, ()):
                    m = self._graph_metrics.get(fid, {})
                    comm = m.get("community_id")
                    if comm is not None:
                        seed_communities[comm] += 1
            if seed_communities:
                total_seeds = sum(seed_communities.values())
                for fid in list(activation.keys()):
                    m = self._graph_metrics.get(fid, {})
                    fact_comm = m.get("community_id")
                    if fact_comm is not None and fact_comm in seed_communities:
                        boost = min(1.0 + 0.15 * (seed_communities[fact_comm] / total_seeds), 1.3)
                        activation[fid] *= boost

        # Extract scores ONLY for the candidate set, normalize to [0, 1]
        candidate_set = set(candidate_fact_ids)
        scored = {fid: activation.get(fid, 0.0) for fid in candidate_set}

        max_score = max(scored.values()) if scored else 0
        if max_score > 0:
            scored = {fid: sc / max_score for fid, sc in scored.items()}

        return scored

    def _suppress_contradictions(
        self, activation: dict[str, float], profile_id: str,
    ) -> None:
        """P3: Penalize older fact in contradiction pairs, heavy-penalize superseded.

        Uses graph_edges (edge_type CHECK includes 'contradiction', 'supersedes').
        """
        candidate_ids = list(activation.keys())
        if not candidate_ids:
            return
        try:
            placeholders = ",".join("?" * len(candidate_ids))
            sql = (
                "SELECT source_id, target_id, edge_type FROM graph_edges "
                "WHERE profile_id = ? AND edge_type IN ('contradiction', 'supersedes') "
                "AND (source_id IN (" + placeholders + ") "
                "OR target_id IN (" + placeholders + "))"
            )
            rows = self._db.execute(sql, (profile_id, *candidate_ids, *candidate_ids))
            edges = [dict(r) for r in rows]
            if not edges:
                return

            # Batch load created_at for involved facts
            involved = set()
            for e in edges:
                involved.add(e["source_id"])
                involved.add(e["target_id"])
            involved = involved & set(candidate_ids)
            if not involved:
                return
            ph2 = ",".join("?" * len(involved))
            ts_rows = self._db.execute(
                "SELECT fact_id, created_at FROM atomic_facts "
                "WHERE fact_id IN (" + ph2 + ") AND profile_id = ?",
                (*involved, profile_id),
            )
            ts_map = {dict(r)["fact_id"]: dict(r).get("created_at", "") for r in ts_rows}

            for e in edges:
                src, tgt, etype = e["source_id"], e["target_id"], e["edge_type"]
                if etype == "supersedes" and src in activation:
                    activation[src] *= 0.3  # Heavy penalty: this fact was replaced
                elif etype == "contradiction":
                    src_ts = ts_map.get(src, "")
                    tgt_ts = ts_map.get(tgt, "")
                    if src_ts and tgt_ts:
                        older = src if src_ts < tgt_ts else tgt
                        if older in activation:
                            activation[older] *= 0.5
        except Exception as exc:
            logger.debug("Contradiction suppression failed: %s", exc)

    def _resolve_entities(self, raw: list[str], profile_id: str) -> list[str]:
        """Resolve raw names to canonical entity IDs."""
        ids: list[str] = []
        seen: set[str] = set()
        if self._resolver is not None:
            for eid in self._resolver.resolve(raw, profile_id).values():
                if eid not in seen:
                    seen.add(eid)
                    ids.append(eid)
        else:
            for name in raw:
                ent = self._db.get_entity_by_name(name, profile_id)
                if ent and ent.entity_id not in seen:
                    seen.add(ent.entity_id)
                    ids.append(ent.entity_id)
        return ids

    def _discover_entities(
        self, fact_ids: set[str], profile_id: str, visited: set[str],
    ) -> list[str]:
        """Find new canonical entity IDs referenced by a set of facts."""
        new: list[str] = []
        seen = set(visited)
        for fid in fact_ids:
            rows = self._db.execute(
                "SELECT canonical_entities_json FROM atomic_facts WHERE fact_id = ?", (fid,),
            )
            if not rows:
                continue
            raw = dict(rows[0]).get("canonical_entities_json")
            if not raw:
                continue
            try:
                for eid in json.loads(raw):
                    if eid not in seen:
                        seen.add(eid)
                        new.append(eid)
            except (ValueError, TypeError):
                continue
        return new

    # v3.4.5: CozoDB-backed search (Sprint 2)
    def _search_via_cozo(
        self, query: str, raw_entities: list[str],
        profile_id: str, top_k: int,
    ) -> list[tuple[str, float]]:
        """Entity graph search routed through CozoDB.

        Uses CozoDB for spreading activation — avoids loading
        the full adjacency graph into memory.
        Falls back to in-memory adjacency if CozoDB fails.
        """
        if not raw_entities:
            return []

        canonical_ids = self._resolve_entities(raw_entities, profile_id)
        if not canonical_ids:
            return []

        try:
            # Use CozoDB for spreading activation
            scored = self._cozo.spreading_activation(
                canonical_ids,
                depth=self._max_hops,
                decay=self._decay,
                top_k=top_k * 2,  # Fetch extra for filtering
            )

            # Map entity scores to fact scores
            fact_scores: list[tuple[str, float]] = []
            for entity_id, score in scored:
                facts = self._db.get_facts_by_entity(entity_id, profile_id)
                for fact in facts:
                    fact_scores.append((fact.fact_id, score))

            # Sort and return top_k
            fact_scores.sort(key=lambda x: x[1], reverse=True)
            return fact_scores[:top_k]

        except Exception:
            # Fallback to in-memory adjacency (existing code path)
            return []
