from __future__ import annotations

"""Deterministic semantic graph enrichment.

This module adds a very lightweight, *non-LLM* semantic layer on top of the
structural graph:

- One summary node per section (concatenation of child paragraph texts).
- Very naive entity nodes by keyword matching (e.g. "RAG", "GraphRAG").

The goal is to:

- Provide a concrete shape for semantic nodes (SUMMARY, ENTITY).
- Make it easy to plug in real NLP/LLM components later without changing
  callers.
"""

from dataclasses import replace
from typing import Dict, Iterable, List

from ..core.models import DocumentGraph, GraphEdge, GraphNode, NodeType
from .base import BaseGraphBuilder


class SemanticGraphBuilder(BaseGraphBuilder):
    """Augment an existing structural graph with summary/entity nodes.

    This builder expects a graph produced by :class:`StructuralGraphBuilder`
    and enriches it in-place.
    """

    def supports(self, tree) -> bool:  # type: ignore[override]
        # For now we always allow semantic enrichment; callers can opt-out by
        # not including this builder in the pipeline.
        return True

    def build(self, tree) -> DocumentGraph:  # type: ignore[override]
        # NOTE: For the v0 implementation we reconstruct a structural graph
        # using the existing builder and then enrich it here. This keeps the
        # public entry point simple for callers.
        from .structure import StructuralGraphBuilder

        base_graph = StructuralGraphBuilder().build(tree)
        return enrich_graph_with_semantics(base_graph)


def enrich_graph_with_semantics(graph: DocumentGraph) -> DocumentGraph:
    """Return a new :class:`DocumentGraph` with semantic nodes added.

    - For each SECTION node with paragraph children, add a SUMMARY node and
      connect it via ``summary_of`` edges.
    - For each paragraph, detect a few hard-coded keywords and create
      corresponding ENTITY nodes, connected via ``mentions`` edges.
    """

    nodes_by_id: Dict[str, GraphNode] = {n.id: n for n in graph.nodes}
    new_nodes: List[GraphNode] = []
    new_edges: List[GraphEdge] = []

    # Build a quick adjacency from structural edges.
    children_by_parent: Dict[str, List[str]] = {}
    for e in graph.edges:
        if e.edge_type == "parent_of":
            children_by_parent.setdefault(e.source, []).append(e.target)

    # --- Summary nodes per SECTION ---
    for node in graph.nodes:
        if node.node_type is not NodeType.SECTION:
            continue

        child_ids = children_by_parent.get(node.id, [])
        paragraph_texts: List[str] = []
        for cid in child_ids:
            child = nodes_by_id.get(cid)
            if child is None:
                continue
            if child.node_type is NodeType.PARAGRAPH and child.label:
                paragraph_texts.append(child.label.strip())

        if not paragraph_texts:
            continue

        summary_text = "\n\n".join(paragraph_texts)
        summary_id = f"{node.id}#summary"
        summary_node = GraphNode(
            id=summary_id,
            node_type=NodeType.SUMMARY,
            label=summary_text,
            metadata={"source_section": node.id},
        )
        new_nodes.append(summary_node)
        new_edges.append(
            GraphEdge(
                source=summary_id,
                target=node.id,
                edge_type="summary_of",
                metadata={},
            )
        )

    # --- Naive entity nodes via keyword matching ---
    KEYWORDS = ["rag", "graphrag", "agent", "summary"]

    for node in graph.nodes:
        if node.node_type is not NodeType.PARAGRAPH or not node.label:
            continue

        text_lower = node.label.lower()

        for kw in KEYWORDS:
            if kw not in text_lower:
                continue

            entity_id = f"{node.id}#ent:{kw}"
            entity_node = GraphNode(
                id=entity_id,
                node_type=NodeType.ENTITY,
                label=kw,
                metadata={"keyword": kw},
            )
            new_nodes.append(entity_node)
            new_edges.append(
                GraphEdge(
                    source=node.id,
                    target=entity_id,
                    edge_type="mentions",
                    metadata={},
                )
            )

    if not new_nodes and not new_edges:
        return graph

    enriched = DocumentGraph(
        nodes=list(graph.nodes) + new_nodes,
        edges=list(graph.edges) + new_edges,
        metadata=dict(graph.metadata or {}),
    )
    enriched.metadata.setdefault("semantic", {})["enabled"] = True
    return enriched

