"""
graph_extractor.py
-------------------
Automatic graph extraction from text for HIEROMEM.

Enables HIEROMEM to work for RAG without manual graph construction.
Extracts entities and relationships from text automatically.

Author: Antigravity
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Set, Tuple
import networkx as nx


# Stopwords to filter from entity extraction
STOPWORDS = {
    "who", "what", "where", "when", "why", "how", "which",
    "the", "a", "an", "is", "are", "was", "were", "be", "been",
    "and", "or", "but", "if", "then", "else", "for", "of", "to",
    "in", "on", "at", "by", "with", "from", "as", "into",
    "this", "that", "these", "those", "it", "its",
    "did", "does", "do", "has", "have", "had",
}


@dataclass
class ExtractedEntity:
    """A named entity extracted from text."""
    text: str
    label: str  # PERSON, ORG, GPE, etc.
    start: int
    end: int


@dataclass
class ExtractedRelation:
    """A relation between two entities."""
    source: str
    target: str
    relation: str
    confidence: float = 1.0


@dataclass
class ExtractionResult:
    """Result of graph extraction from text."""
    entities: List[ExtractedEntity]
    relations: List[ExtractedRelation]
    graph: nx.DiGraph

    @property
    def entity_texts(self) -> Set[str]:
        return {e.text.lower() for e in self.entities}


class GraphExtractor(ABC):
    """Base class for graph extraction strategies."""

    @abstractmethod
    def extract(self, text: str, context: Optional[Dict[str, Any]] = None) -> ExtractionResult:
        """Extract entities and relations from text."""
        pass

    def extract_batch(self, texts: List[str]) -> List[ExtractionResult]:
        """Extract from multiple texts."""
        return [self.extract(t) for t in texts]


class RegexGraphExtractor(GraphExtractor):
    """
    Regex-based entity extraction (fallback).
    Fast but lower quality - extracts capitalized phrases.
    """

    def __init__(self):
        # Patterns for different entity types
        self.patterns = {
            "PERSON": r'\b[A-Z][a-z]+(?:\s+[A-Z][a-z]+)+\b',
            "ORG": (
                r'\b(?:The\s+)?[A-Z][a-zA-Z]*(?:\s+[A-Z][a-zA-Z]*)*'
                r'(?:\s+(?:Inc|Corp|Ltd|LLC|University|College|Institute|Company))?\b'
            ),
            "GPE": r'\b[A-Z][a-z]+(?:,\s*[A-Z][a-z]+)?\b',
        }

    def extract(self, text: str, context: Optional[Dict[str, Any]] = None) -> ExtractionResult:
        entities = []
        seen = set()

        # Extract capitalized phrases (2-4 words)
        words = text.split()
        i = 0
        while i < len(words):
            for length in [4, 3, 2, 1]:
                if i + length <= len(words):
                    phrase = " ".join(words[i:i+length])
                    phrase_words = phrase.split()
                    # Check if all significant words start with capital
                    if all(w[0].isupper() for w in phrase_words if len(w) > 2):
                        clean = phrase.strip(".,!?;:'\"()[]")
                        if len(clean) > 2 and clean.lower() not in seen:
                            seen.add(clean.lower())
                            entities.append(ExtractedEntity(
                                text=clean,
                                label="ENTITY",
                                start=text.find(clean),
                                end=text.find(clean) + len(clean)
                            ))
                            i += length - 1
                            break
            i += 1

        # Filter stopwords
        entities = [e for e in entities if e.text.lower() not in STOPWORDS]

        # Build co-occurrence graph
        G = nx.DiGraph()
        for ent in entities:
            G.add_node(ent.text.lower(), label=ent.label)

        # Add co-occurrence edges between entities in same text
        relations = []
        entity_texts = [e.text.lower() for e in entities]
        for i, e1 in enumerate(entity_texts):
            for e2 in entity_texts[i+1:]:
                if e1 != e2:
                    G.add_edge(e1, e2, relation="co_occurs")
                    relations.append(ExtractedRelation(e1, e2, "co_occurs"))

        return ExtractionResult(entities=entities, relations=relations, graph=G)


class SpacyGraphExtractor(GraphExtractor):
    """
    SpaCy-based NER extraction (recommended).
    Higher quality entity recognition.
    """

    def __init__(self, model_name: str = "en_core_web_sm"):
        self.model_name = model_name
        self._nlp = None

    @property
    def nlp(self):
        if self._nlp is None:
            try:
                import spacy
                self._nlp = spacy.load(self.model_name)
            except ImportError:
                raise ImportError(
                    "SpaCy not installed. Run: pip install spacy && "
                    "python -m spacy download en_core_web_sm"
                )
            except OSError:
                raise OSError(
                    f"SpaCy model '{self.model_name}' not found. "
                    f"Run: python -m spacy download {self.model_name}"
                )
        return self._nlp

    def extract(self, text: str, context: Optional[Dict[str, Any]] = None) -> ExtractionResult:
        doc = self.nlp(text)

        entities = []
        for ent in doc.ents:
            if ent.label_ in {"PERSON", "ORG", "GPE", "LOC", "WORK_OF_ART", "EVENT", "PRODUCT", "NORP"}:
                entities.append(ExtractedEntity(
                    text=ent.text,
                    label=ent.label_,
                    start=ent.start_char,
                    end=ent.end_char
                ))

        # Build graph
        G = nx.DiGraph()
        for ent in entities:
            G.add_node(ent.text.lower(), label=ent.label)

        # Add noun chunk entities as fallback
        for chunk in doc.noun_chunks:
            if chunk.root.pos_ == "PROPN":
                key = chunk.text.lower()
                if key not in G:
                    G.add_node(key, label="NOUN_CHUNK")
                    entities.append(ExtractedEntity(
                        text=chunk.text,
                        label="NOUN_CHUNK",
                        start=chunk.start_char,
                        end=chunk.end_char
                    ))

        # Add co-occurrence edges
        relations = []
        entity_texts = list(set(e.text.lower() for e in entities))
        for i, e1 in enumerate(entity_texts):
            for e2 in entity_texts[i+1:]:
                if e1 != e2:
                    G.add_edge(e1, e2, relation="co_occurs")
                    relations.append(ExtractedRelation(e1, e2, "co_occurs"))

        # Filter stopwords
        entities = [e for e in entities if e.text.lower() not in STOPWORDS]

        return ExtractionResult(entities=entities, relations=relations, graph=G)


class CrossDocumentLinker:
    """
    Links entities across documents to build a corpus-wide graph.
    This is the key to making HIEROMEM work for RAG.
    """

    def __init__(self):
        self.entity_to_docs: Dict[str, Set[str]] = {}  # entity -> set of doc_ids
        self.doc_to_entities: Dict[str, Set[str]] = {}  # doc_id -> set of entities

    def register(self, doc_id: str, entities: Set[str]) -> None:
        """Register entities for a document."""
        self.doc_to_entities[doc_id] = entities
        for ent in entities:
            self.entity_to_docs.setdefault(ent, set()).add(doc_id)

    def get_related_docs(self, doc_id: str, min_shared: int = 1) -> List[Tuple[str, int]]:
        """Get documents related to this one via shared entities."""
        if doc_id not in self.doc_to_entities:
            return []

        my_entities = self.doc_to_entities[doc_id]
        doc_scores = {}

        for ent in my_entities:
            for other_doc in self.entity_to_docs.get(ent, set()):
                if other_doc != doc_id:
                    doc_scores[other_doc] = doc_scores.get(other_doc, 0) + 1

        related = [(doc, score) for doc, score in doc_scores.items() if score >= min_shared]
        return sorted(related, key=lambda x: -x[1])

    def build_cross_doc_edges(self, doc_id: str, limit: int = 20) -> List[Tuple[str, str, Dict]]:
        """Build edges from this doc to related docs."""
        edges = []
        for related_doc, score in self.get_related_docs(doc_id)[:limit]:
            # Find shared entities for edge metadata
            my_ents = self.doc_to_entities.get(doc_id, set())
            their_ents = self.doc_to_entities.get(related_doc, set())
            shared = my_ents & their_ents

            edges.append((doc_id, related_doc, {
                "relation": "shares_entities",
                "shared_count": score,
                "shared_entities": list(shared)[:5]  # Limit metadata size
            }))

        return edges


class AutoGraphExtractor:
    """
    Main entry point for automatic graph extraction.
    Combines NER with cross-document linking.
    """

    def __init__(
        self,
        backend: str = "spacy",
        spacy_model: str = "en_core_web_sm",
        enable_cross_doc: bool = True,
        max_cross_doc_edges: int = 20
    ):
        self.backend = backend
        self.enable_cross_doc = enable_cross_doc
        self.max_cross_doc_edges = max_cross_doc_edges

        if backend == "spacy":
            self._extractor = SpacyGraphExtractor(spacy_model)
        else:
            self._extractor = RegexGraphExtractor()

        self._linker = CrossDocumentLinker() if enable_cross_doc else None

    def extract(self, text: str, doc_id: Optional[str] = None) -> nx.DiGraph:
        """
        Extract a graph from text.

        If doc_id is provided and cross-doc linking is enabled,
        also adds edges to related documents.
        """
        result = self._extractor.extract(text)
        G = result.graph.copy()

        # Register entities for cross-doc linking
        if self._linker and doc_id:
            self._linker.register(doc_id, result.entity_texts)

            # Add cross-document edges
            for src, dst, data in self._linker.build_cross_doc_edges(
                doc_id, limit=self.max_cross_doc_edges
            ):
                G.add_edge(src, dst, **data)

        return G

    def extract_entities(self, text: str) -> Set[str]:
        """Extract just entity strings (for query graphs)."""
        result = self._extractor.extract(text)
        return result.entity_texts

    @property
    def linker(self) -> Optional[CrossDocumentLinker]:
        """Access the cross-document linker."""
        return self._linker
