"""
cogniq.graph — Lightweight knowledge graph for causal + semantic reasoning.

Extracts entities and relations from text, enables causal path traversal
to augment pure embedding retrieval (the "Semantic DNA" layer).
"""

from __future__ import annotations
import re
from collections import defaultdict, deque
from dataclasses import dataclass, field
from typing import List, Dict, Tuple, Set, Optional, Any


@dataclass
class Entity:
    name: str
    type: str = "concept"
    mentions: int = 1
    sources: List[str] = field(default_factory=list)

    def __hash__(self):
        return hash(self.name.lower())

    def __eq__(self, other):
        return self.name.lower() == other.name.lower()


@dataclass
class Relation:
    source: str   # entity name
    target: str
    rel_type: str  # "causes", "requires", "enables", "part_of", etc.
    weight: float = 1.0
    evidence: str = ""


class KnowledgeGraph:
    """
    Auto-builds a knowledge graph from ingested documents.

    Features:
    ─ Entity extraction (noun phrases, capitalized terms, domain terms)
    ─ Relation extraction (causal, hierarchical, definitional patterns)
    ─ BFS path finding for causal chains
    ─ Graph-augmented retrieval scoring
    ─ Custom entity/relation injection
    """

    # Causal patterns (regex → relation_type)
    CAUSAL_PATTERNS: List[Tuple[str, str]] = [
        (r"(\w[\w\s]+)\s+causes?\s+(\w[\w\s]+)", "causes"),
        (r"(\w[\w\s]+)\s+leads?\s+to\s+(\w[\w\s]+)", "leads_to"),
        (r"(\w[\w\s]+)\s+results?\s+in\s+(\w[\w\s]+)", "results_in"),
        (r"(\w[\w\s]+)\s+triggers?\s+(\w[\w\s]+)", "triggers"),
        (r"(\w[\w\s]+)\s+requires?\s+(\w[\w\s]+)", "requires"),
        (r"(\w[\w\s]+)\s+enables?\s+(\w[\w\s]+)", "enables"),
        (r"(\w[\w\s]+)\s+prevents?\s+(\w[\w\s]+)", "prevents"),
        (r"(\w[\w\s]+)\s+(?:is a|are a|is an)\s+(\w[\w\s]+)", "is_a"),
        (r"(\w[\w\s]+)\s+(?:include|includes|such as)\s+(\w[\w\s]+)", "includes"),
        (r"(\w[\w\s]+)\s+(?:must|shall|should)\s+(?:maintain|hold|have)\s+(\w[\w\s]+)", "obligates"),
        (r"failure to\s+(\w[\w\s]+)\s+attracts?\s+(\w[\w\s]+)", "violation_triggers"),
        (r"(\w[\w\s]+)\s+(?:impacts?|affects?)\s+(\w[\w\s]+)", "impacts"),
    ]

    def __init__(self):
        self._entities: Dict[str, Entity] = {}   # name.lower() → Entity
        self._adj: Dict[str, List[Relation]] = defaultdict(list)  # out-edges
        self._rev: Dict[str, List[Relation]] = defaultdict(list)  # in-edges (for reverse traversal)
        self._compiled = [
            (re.compile(pat, re.IGNORECASE), rel)
            for pat, rel in self.CAUSAL_PATTERNS
        ]

    # ─── Ingest ───────────────────────────────────────────────

    def ingest(self, text: str, source: Optional[str] = None):
        """Extract entities and relations from text."""
        entities = self._extract_entities(text)
        for e in entities:
            self._add_entity(e, source)

        relations = self._extract_relations(text)
        for r in relations:
            self._add_relation(r)

    def _extract_entities(self, text: str) -> List[Entity]:
        entities = []
        # Capitalized multi-word terms (likely proper nouns / concepts)
        caps = re.findall(r'\b([A-Z][A-Za-z]+(?:\s+[A-Z][A-Za-z]+)*)\b', text)
        for cap in caps:
            if len(cap) > 2 and cap.lower() not in {"The", "A", "An", "This", "That"}:
                entities.append(Entity(name=cap, type="term"))

        # Acronyms (LCR, HQLA, RBI, etc.)
        acronyms = re.findall(r'\b([A-Z]{2,6})\b', text)
        for acro in acronyms:
            entities.append(Entity(name=acro, type="acronym"))

        return entities

    def _extract_relations(self, text: str) -> List[Relation]:
        relations = []
        # Sentence-by-sentence extraction
        sentences = re.split(r'[.!?]\s+', text)
        for sent in sentences:
            for pattern, rel_type in self._compiled:
                for match in pattern.finditer(sent):
                    subj = match.group(1).strip()[:50]
                    obj = match.group(2).strip()[:50]
                    if len(subj) > 2 and len(obj) > 2:
                        relations.append(Relation(
                            source=subj, target=obj,
                            rel_type=rel_type, evidence=sent[:100]
                        ))
        return relations

    def _add_entity(self, entity: Entity, source: Optional[str]):
        key = entity.name.lower()
        if key in self._entities:
            self._entities[key].mentions += 1
            if source:
                self._entities[key].sources.append(source)
        else:
            if source:
                entity.sources.append(source)
            self._entities[key] = entity

    def _add_relation(self, rel: Relation):
        self._adj[rel.source.lower()].append(rel)
        self._rev[rel.target.lower()].append(rel)

    # ─── Custom injection ─────────────────────────────────────

    def add_entity(self, name: str, entity_type: str = "custom"):
        """Manually add an entity."""
        self._add_entity(Entity(name=name, type=entity_type), source="custom")

    def add_relation(self, source: str, target: str, rel_type: str, weight: float = 1.0):
        """Manually add a relation."""
        self._add_relation(Relation(source=source, target=target,
                                     rel_type=rel_type, weight=weight))

    # ─── Traversal ────────────────────────────────────────────

    def find_paths(self, start: str, end: str, max_depth: int = 4) -> List[List[Relation]]:
        """BFS to find all causal paths from start to end."""
        start, end = start.lower(), end.lower()
        if start not in self._adj:
            return []

        queue = deque([(start, [])])
        visited: Set[str] = {start}
        paths = []

        while queue:
            node, path = queue.popleft()
            if len(path) >= max_depth:
                continue
            for rel in self._adj.get(node, []):
                next_node = rel.target.lower()
                new_path = path + [rel]
                if next_node == end:
                    paths.append(new_path)
                elif next_node not in visited:
                    visited.add(next_node)
                    queue.append((next_node, new_path))

        return paths

    def neighbors(self, entity: str, direction: str = "out") -> List[Relation]:
        """Get direct neighbors (outgoing or incoming relations)."""
        key = entity.lower()
        if direction == "out":
            return self._adj.get(key, [])
        return self._rev.get(key, [])

    def subgraph(self, entities: List[str], depth: int = 2) -> "KnowledgeGraph":
        """Extract a subgraph around given entities."""
        sg = KnowledgeGraph()
        visited: Set[str] = set()
        queue = deque([(e.lower(), 0) for e in entities])

        while queue:
            node, d = queue.popleft()
            if node in visited or d > depth:
                continue
            visited.add(node)
            if node in self._entities:
                sg._entities[node] = self._entities[node]
            for rel in self._adj.get(node, []):
                sg._add_relation(rel)
                if d < depth:
                    queue.append((rel.target.lower(), d + 1))

        return sg

    # ─── Scoring ──────────────────────────────────────────────

    def score_relevance(self, query: str, doc_text: str) -> float:
        """
        Graph-based relevance boost:
        Count shared entities between query and doc, weighted by
        their centrality (mention count) in the graph.
        """
        q_ents = {e.name.lower() for e in self._extract_entities(query)}
        d_ents = {e.name.lower() for e in self._extract_entities(doc_text)}
        shared = q_ents & d_ents
        if not shared:
            return 0.0
        score = sum(
            self._entities.get(e, Entity(e)).mentions
            for e in shared
        )
        # Normalize by max possible
        max_score = len(q_ents) * max(
            (self._entities.get(e, Entity(e)).mentions for e in q_ents), default=1
        )
        return min(score / max(max_score, 1), 1.0)

    # ─── Stats ────────────────────────────────────────────────

    def stats(self) -> Dict[str, Any]:
        total_edges = sum(len(v) for v in self._adj.values())
        return {
            "entities": len(self._entities),
            "relations": total_edges,
            "top_entities": sorted(
                self._entities.values(), key=lambda e: e.mentions, reverse=True
            )[:10],
        }

    def __repr__(self):
        s = self.stats()
        return f"KnowledgeGraph(entities={s['entities']}, relations={s['relations']})"
