#!/usr/bin/env python3
"""GraphRAG knowledge graph extraction demo — entity recognition + relation extraction.

Demonstrates a minimal pipeline: NER → relation extraction → entity resolution → graph construction.
Uses heuristic extraction (no LLM needed) to illustrate the pipeline shape.
"""

from __future__ import annotations

import re
from collections import defaultdict
from dataclasses import dataclass, field


@dataclass
class Entity:
    id: str
    name: str
    entity_type: str
    aliases: set[str] = field(default_factory=set)


@dataclass
class Relation:
    source: str
    target: str
    relation_type: str


class KnowledgeGraph:
    def __init__(self):
        self.entities: dict[str, Entity] = {}
        self.relations: list[Relation] = []

    def add_entity(self, entity: Entity):
        self.entities[entity.id] = entity

    def add_relation(self, rel: Relation):
        self.relations.append(rel)

    def resolve_aliases(self):
        """Merge entities with overlapping aliases (simple heuristic)."""
        merged: dict[str, Entity] = {}
        canonical: dict[str, str] = {}  # alias → canonical id

        for eid, entity in self.entities.items():
            # Check if any alias already seen
            existing = None
            for alias in entity.aliases | {entity.name}:
                if alias in canonical:
                    existing = canonical[alias]
                    break
            if existing:
                # Merge into existing
                merged[existing].aliases |= entity.aliases | {entity.name}
            else:
                merged[eid] = entity
                for alias in entity.aliases | {entity.name}:
                    canonical[alias] = eid

        self.entities = merged

    def query_neighbors(self, entity_name: str, hops: int = 1) -> list[str]:
        """Find entities related to the given entity name."""
        entity_id = None
        for eid, ent in self.entities.items():
            if entity_name in ent.aliases or entity_name == ent.name:
                entity_id = eid
                break
        if not entity_id:
            return []

        related: set[str] = set()
        for rel in self.relations:
            if rel.source == entity_id:
                target = self.entities.get(rel.target)
                if target:
                    related.add(f"{target.name} --[{rel.relation_type}]→ {entity_name}")
            if rel.target == entity_id:
                source = self.entities.get(rel.source)
                if source:
                    related.add(f"{entity_name} --[{rel.relation_type}]→ {source.name}")
        return list(related)


# ── Heuristic extraction (replace with LLM in production) ───────────────

def extract_entities_and_relations(texts: list[str]) -> KnowledgeGraph:
    """Extract entities and relations from text via simple regex patterns."""
    kg = KnowledgeGraph()

    for text in texts:
        # Simple entity patterns
        for match in re.finditer(r"(Service|System|Team|Product):\s*(\w+)", text):
            entity_type, name = match.groups()
            eid = f"{entity_type.lower()}_{name.lower()}"
            kg.add_entity(Entity(id=eid, name=name, entity_type=entity_type))

        # Simple relation patterns
        for match in re.finditer(r"(\w+)\s+(depends on|manages|uses|owns)\s+(\w+)", text):
            src, rel_type, tgt = match.groups()
            kg.add_relation(Relation(source=src, target=tgt, relation_type=rel_type))

    kg.resolve_aliases()
    return kg


# ── Demo ────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    docs = [
        "Service: AuthService depends on Service: DatabaseService. Team: Platform uses Service: AuthService.",
        "Service: DatabaseService depends on Service: CacheService. Team: DataEng manages Service: DatabaseService.",
        "Product: Dashboard uses Service: AuthService and depends on Service: CacheService.",
        "The AuthService is also called AuthSvc by the Platform team.",
    ]

    kg = extract_entities_and_relations(docs)

    print(f"Entities ({len(kg.entities)}):")
    for eid, ent in kg.entities.items():
        print(f"  [{ent.entity_type}] {ent.name}")

    print(f"\nRelations ({len(kg.relations)}):")
    for rel in kg.relations:
        print(f"  {rel.source} --[{rel.relation_type}]→ {rel.target}")

    print(f"\nMulti-hop query: neighbors of AuthService:")
    for path in kg.query_neighbors("AuthService"):
        print(f"  {path}")

    print("\n✓ GraphRAG extraction demo complete.")