"""
Farnsworth Knowledge Graph - Entity Relationship Memory

Novel Approaches:
1. Auto-Extraction - Automatic entity and relation extraction
2. Semantic Linking - Embedding-based relationship discovery
3. Temporal Edges - Time-aware relationship tracking
4. Inference Paths - Multi-hop reasoning support
"""

import asyncio
import json
import re
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Optional, Any, Callable
from collections import defaultdict

from loguru import logger

try:
    import networkx as nx
except ImportError:
    nx = None
    logger.warning("NetworkX not installed, knowledge graph features limited")


@dataclass
class Entity:
    """An entity in the knowledge graph."""
    id: str
    name: str
    entity_type: str  # person, concept, tool, file, etc.
    properties: dict = field(default_factory=dict)
    created_at: datetime = field(default_factory=datetime.now)
    last_mentioned: datetime = field(default_factory=datetime.now)
    mention_count: int = 1
    embedding: Optional[list[float]] = None

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "name": self.name,
            "entity_type": self.entity_type,
            "properties": self.properties,
            "created_at": self.created_at.isoformat(),
            "last_mentioned": self.last_mentioned.isoformat(),
            "mention_count": self.mention_count,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "Entity":
        return cls(
            id=data["id"],
            name=data["name"],
            entity_type=data["entity_type"],
            properties=data.get("properties", {}),
            created_at=datetime.fromisoformat(data.get("created_at", datetime.now().isoformat())),
            last_mentioned=datetime.fromisoformat(data.get("last_mentioned", datetime.now().isoformat())),
            mention_count=data.get("mention_count", 1),
        )


@dataclass
class Relationship:
    """A relationship between entities."""
    source_id: str
    target_id: str
    relation_type: str  # uses, is_a, part_of, relates_to, etc.
    weight: float = 1.0
    created_at: datetime = field(default_factory=datetime.now)
    last_updated: datetime = field(default_factory=datetime.now)
    evidence: list[str] = field(default_factory=list)  # Source texts
    bidirectional: bool = False

    def to_dict(self) -> dict:
        return {
            "source_id": self.source_id,
            "target_id": self.target_id,
            "relation_type": self.relation_type,
            "weight": self.weight,
            "created_at": self.created_at.isoformat(),
            "last_updated": self.last_updated.isoformat(),
            "evidence": self.evidence[:5],  # Limit stored evidence
            "bidirectional": self.bidirectional,
        }


@dataclass
class GraphQuery:
    """Query result from knowledge graph."""
    entities: list[Entity]
    relationships: list[Relationship]
    paths: list[list[str]] = field(default_factory=list)
    score: float = 0.0


class KnowledgeGraph:
    """
    NetworkX-based knowledge graph for entity relationships.

    Features:
    - Automatic entity extraction from text
    - Semantic relationship discovery
    - Multi-hop path finding
    - Graph-based retrieval for RAG
    """

    def __init__(
        self,
        data_dir: str = "./data/graph",
        max_nodes: int = 10000,
        auto_link_threshold: float = 0.75,
    ):
        self.data_dir = Path(data_dir)
        self.data_dir.mkdir(parents=True, exist_ok=True)

        self.max_nodes = max_nodes
        self.auto_link_threshold = auto_link_threshold

        # NetworkX graph
        if nx:
            self.graph = nx.DiGraph()
        else:
            self.graph = None

        # Entity storage
        self.entities: dict[str, Entity] = {}
        self.name_to_id: dict[str, str] = {}  # Lowercase name -> entity_id

        # Type indices
        self.entities_by_type: dict[str, set[str]] = defaultdict(set)

        # Embedding function (set externally)
        self.embed_fn: Optional[Callable] = None

        self._lock = asyncio.Lock()

    async def initialize(self):
        """Load existing graph from disk."""
        await self._load_from_disk()
        logger.info(f"Knowledge graph initialized with {len(self.entities)} entities")

    async def add_entity(
        self,
        name: str,
        entity_type: str,
        properties: Optional[dict] = None,
        embedding: Optional[list[float]] = None,
    ) -> Entity:
        """Add or update an entity."""
        async with self._lock:
            # Check for existing entity
            name_lower = name.lower()
            if name_lower in self.name_to_id:
                entity = self.entities[self.name_to_id[name_lower]]
                entity.mention_count += 1
                entity.last_mentioned = datetime.now()
                if properties:
                    entity.properties.update(properties)
                return entity

            # Create new entity
            entity_id = f"e_{len(self.entities)}_{hash(name) % 10000:04d}"

            # Get embedding if function available
            if embedding is None and self.embed_fn:
                embedding = await self._get_embedding(name)

            entity = Entity(
                id=entity_id,
                name=name,
                entity_type=entity_type,
                properties=properties or {},
                embedding=embedding,
            )

            # Check capacity
            if len(self.entities) >= self.max_nodes:
                await self._prune_least_important()

            # Add to storage
            self.entities[entity_id] = entity
            self.name_to_id[name_lower] = entity_id
            self.entities_by_type[entity_type].add(entity_id)

            # Add to graph
            if self.graph is not None:
                self.graph.add_node(entity_id, **entity.to_dict())

            # Auto-link to similar entities
            if embedding:
                await self._auto_link_entity(entity)

            return entity

    async def add_relationship(
        self,
        source: str,  # Entity name or ID
        target: str,  # Entity name or ID
        relation_type: str,
        weight: float = 1.0,
        evidence: Optional[str] = None,
        bidirectional: bool = False,
    ) -> Relationship:
        """Add a relationship between entities."""
        async with self._lock:
            # Resolve entity IDs
            source_id = self._resolve_entity(source)
            target_id = self._resolve_entity(target)

            if not source_id or not target_id:
                raise ValueError(f"Entity not found: {source if not source_id else target}")

            # Check for existing relationship
            if self.graph is not None:
                if self.graph.has_edge(source_id, target_id):
                    # Update existing
                    edge_data = self.graph.edges[source_id, target_id]
                    edge_data["weight"] = edge_data.get("weight", 1.0) + weight * 0.1
                    edge_data["last_updated"] = datetime.now().isoformat()
                    if evidence:
                        edge_data.setdefault("evidence", []).append(evidence)

                    return Relationship(
                        source_id=source_id,
                        target_id=target_id,
                        relation_type=edge_data.get("relation_type", relation_type),
                        weight=edge_data["weight"],
                    )

            # Create relationship
            rel = Relationship(
                source_id=source_id,
                target_id=target_id,
                relation_type=relation_type,
                weight=weight,
                evidence=[evidence] if evidence else [],
                bidirectional=bidirectional,
            )

            # Add to graph
            if self.graph is not None:
                self.graph.add_edge(source_id, target_id, **rel.to_dict())
                if bidirectional:
                    self.graph.add_edge(target_id, source_id, **rel.to_dict())

            return rel

    def _resolve_entity(self, name_or_id: str) -> Optional[str]:
        """Resolve entity name or ID to ID."""
        if name_or_id in self.entities:
            return name_or_id
        name_lower = name_or_id.lower()
        return self.name_to_id.get(name_lower)

    async def _auto_link_entity(self, entity: Entity):
        """Automatically link entity to similar entities."""
        if not entity.embedding or not self.embed_fn:
            return

        import numpy as np

        entity_vec = np.array(entity.embedding)

        for other_id, other in self.entities.items():
            if other_id == entity.id or not other.embedding:
                continue

            other_vec = np.array(other.embedding)
            similarity = np.dot(entity_vec, other_vec) / (
                np.linalg.norm(entity_vec) * np.linalg.norm(other_vec) + 1e-8
            )

            if similarity >= self.auto_link_threshold:
                await self.add_relationship(
                    entity.id,
                    other_id,
                    "relates_to",
                    weight=float(similarity),
                    bidirectional=True,
                )

    async def extract_entities_from_text(
        self,
        text: str,
        auto_add: bool = True,
    ) -> list[Entity]:
        """
        Extract entities from text using patterns.

        This is a simple rule-based extractor.
        Could be enhanced with NER models.
        """
        entities = []

        # Extract capitalized phrases (potential proper nouns)
        capitalized = re.findall(r'\b[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*\b', text)
        for phrase in set(capitalized):
            if len(phrase) > 2:
                entity = await self.add_entity(phrase, "concept") if auto_add else Entity(
                    id="", name=phrase, entity_type="concept"
                )
                entities.append(entity)

        # Extract code identifiers (CamelCase or snake_case)
        code_patterns = re.findall(r'\b[a-z]+(?:_[a-z]+)+\b|\b[A-Z][a-z]+(?:[A-Z][a-z]+)+\b', text)
        for pattern in set(code_patterns):
            entity = await self.add_entity(pattern, "code") if auto_add else Entity(
                id="", name=pattern, entity_type="code"
            )
            entities.append(entity)

        # Extract file paths
        file_patterns = re.findall(r'[\w./\\]+\.[a-zA-Z]{1,4}\b', text)
        for fp in set(file_patterns):
            entity = await self.add_entity(fp, "file") if auto_add else Entity(
                id="", name=fp, entity_type="file"
            )
            entities.append(entity)

        # Extract URLs
        url_patterns = re.findall(r'https?://[^\s<>"{}|\\^`\[\]]+', text)
        for url in set(url_patterns):
            entity = await self.add_entity(url, "url") if auto_add else Entity(
                id="", name=url, entity_type="url"
            )
            entities.append(entity)

        return entities

    async def extract_relationships_from_text(
        self,
        text: str,
        entities: list[Entity],
    ) -> list[Relationship]:
        """
        Extract relationships between entities mentioned in text.

        Uses co-occurrence and pattern matching.
        """
        relationships = []

        # Co-occurrence based relationships
        entity_names = [e.name.lower() for e in entities]
        text_lower = text.lower()

        for i, e1 in enumerate(entities):
            for e2 in entities[i+1:]:
                # Check if both entities are mentioned
                if e1.name.lower() in text_lower and e2.name.lower() in text_lower:
                    # Determine relationship type from context
                    rel_type = self._infer_relationship_type(text, e1.name, e2.name)
                    rel = await self.add_relationship(
                        e1.id, e2.id, rel_type,
                        evidence=text[:200],
                        bidirectional=True,
                    )
                    relationships.append(rel)

        return relationships

    def _infer_relationship_type(self, text: str, entity1: str, entity2: str) -> str:
        """Infer relationship type from context."""
        text_lower = text.lower()

        # Check for explicit relationship patterns
        patterns = {
            "uses": [f"{entity1.lower()} uses {entity2.lower()}", f"using {entity2.lower()}"],
            "is_a": [f"{entity1.lower()} is a {entity2.lower()}", f"{entity1.lower()} is an {entity2.lower()}"],
            "part_of": [f"{entity1.lower()} in {entity2.lower()}", f"part of {entity2.lower()}"],
            "creates": [f"{entity1.lower()} creates {entity2.lower()}", f"generate {entity2.lower()}"],
            "depends_on": [f"{entity1.lower()} depends on {entity2.lower()}", f"requires {entity2.lower()}"],
        }

        for rel_type, phrases in patterns.items():
            if any(phrase in text_lower for phrase in phrases):
                return rel_type

        return "relates_to"

    async def query(
        self,
        query: str,
        max_entities: int = 10,
        max_hops: int = 2,
    ) -> GraphQuery:
        """
        Query the knowledge graph.

        Finds relevant entities and their relationships.
        """
        async with self._lock:
            # Extract entities from query
            query_entities = await self.extract_entities_from_text(query, auto_add=False)

            # Find matching entities
            matched_entities = []
            for qe in query_entities:
                entity_id = self._resolve_entity(qe.name)
                if entity_id:
                    matched_entities.append(self.entities[entity_id])

            # Also search by substring
            query_lower = query.lower()
            for name, entity_id in self.name_to_id.items():
                if name in query_lower or query_lower in name:
                    if entity_id not in [e.id for e in matched_entities]:
                        matched_entities.append(self.entities[entity_id])
                        if len(matched_entities) >= max_entities:
                            break

            # Get relationships between matched entities
            relationships = []
            entity_ids = {e.id for e in matched_entities}

            if self.graph is not None:
                for e_id in entity_ids:
                    # Outgoing edges
                    for _, target, data in self.graph.out_edges(e_id, data=True):
                        if target in entity_ids:
                            relationships.append(Relationship(
                                source_id=e_id,
                                target_id=target,
                                relation_type=data.get("relation_type", "relates_to"),
                                weight=data.get("weight", 1.0),
                            ))

                # Find paths between entities
                paths = []
                entity_list = list(entity_ids)
                for i in range(len(entity_list)):
                    for j in range(i + 1, len(entity_list)):
                        try:
                            path = nx.shortest_path(
                                self.graph,
                                entity_list[i],
                                entity_list[j],
                            )
                            if len(path) <= max_hops + 1:
                                paths.append(path)
                        except (nx.NetworkXNoPath, nx.NodeNotFound):
                            pass

            return GraphQuery(
                entities=matched_entities[:max_entities],
                relationships=relationships,
                paths=paths[:10],
                score=len(matched_entities) / max(1, len(query_entities)),
            )

    async def get_neighbors(
        self,
        entity_name_or_id: str,
        max_hops: int = 1,
        relation_filter: Optional[list[str]] = None,
    ) -> list[Entity]:
        """Get neighboring entities."""
        entity_id = self._resolve_entity(entity_name_or_id)
        if not entity_id or self.graph is None:
            return []

        neighbors = set()

        # BFS for multi-hop
        current_level = {entity_id}
        for _ in range(max_hops):
            next_level = set()
            for node in current_level:
                for neighbor in self.graph.neighbors(node):
                    if relation_filter:
                        edge_data = self.graph.edges[node, neighbor]
                        if edge_data.get("relation_type") not in relation_filter:
                            continue
                    if neighbor != entity_id:
                        neighbors.add(neighbor)
                        next_level.add(neighbor)
            current_level = next_level

        return [self.entities[n] for n in neighbors if n in self.entities]

    def get_entity(self, name_or_id: str) -> Optional[Entity]:
        """Get an entity by name or ID."""
        entity_id = self._resolve_entity(name_or_id)
        return self.entities.get(entity_id) if entity_id else None

    async def _prune_least_important(self):
        """Remove least important entities when at capacity."""
        if not self.entities:
            return

        # Score entities
        def importance(e: Entity) -> float:
            recency = (datetime.now() - e.last_mentioned).total_seconds() / 86400
            return e.mention_count / (1 + recency)

        sorted_entities = sorted(self.entities.values(), key=importance)

        # Remove bottom 10%
        to_remove = sorted_entities[:max(1, len(sorted_entities) // 10)]
        for entity in to_remove:
            del self.entities[entity.id]
            self.name_to_id.pop(entity.name.lower(), None)
            self.entities_by_type[entity.entity_type].discard(entity.id)
            if self.graph is not None and entity.id in self.graph:
                self.graph.remove_node(entity.id)

    async def _get_embedding(self, text: str) -> Optional[list[float]]:
        """Get embedding for text."""
        if not self.embed_fn:
            return None
        try:
            if asyncio.iscoroutinefunction(self.embed_fn):
                return await self.embed_fn(text)
            return self.embed_fn(text)
        except Exception as e:
            logger.error(f"Embedding error: {e}")
            return None

    async def save(self):
        """Save graph to disk."""
        # Save entities
        entities_file = self.data_dir / "entities.json"
        entities_data = {eid: e.to_dict() for eid, e in self.entities.items()}
        entities_file.write_text(json.dumps(entities_data, indent=2), encoding='utf-8')

        # Save graph edges
        if self.graph is not None:
            edges_file = self.data_dir / "edges.json"
            edges_data = [
                {"source": u, "target": v, **d}
                for u, v, d in self.graph.edges(data=True)
            ]
            edges_file.write_text(json.dumps(edges_data, indent=2), encoding='utf-8')

    async def _load_from_disk(self):
        """Load graph from disk."""
        entities_file = self.data_dir / "entities.json"
        if entities_file.exists():
            entities_data = json.loads(entities_file.read_text(encoding='utf-8'))
            for eid, data in entities_data.items():
                entity = Entity.from_dict(data)
                self.entities[eid] = entity
                self.name_to_id[entity.name.lower()] = eid
                self.entities_by_type[entity.entity_type].add(eid)
                if self.graph is not None:
                    self.graph.add_node(eid, **entity.to_dict())

        edges_file = self.data_dir / "edges.json"
        if edges_file.exists() and self.graph is not None:
            edges_data = json.loads(edges_file.read_text(encoding='utf-8'))
            for edge in edges_data:
                self.graph.add_edge(edge["source"], edge["target"], **{
                    k: v for k, v in edge.items() if k not in ("source", "target")
                })

    def get_stats(self) -> dict:
        """Get graph statistics."""
        stats = {
            "total_entities": len(self.entities),
            "entities_by_type": {t: len(ids) for t, ids in self.entities_by_type.items()},
            "max_nodes": self.max_nodes,
        }

        if self.graph is not None:
            stats["total_edges"] = self.graph.number_of_edges()
            stats["avg_degree"] = sum(dict(self.graph.degree()).values()) / max(1, len(self.entities))

        return stats
