"""
Context Graph Implementation

This module provides a synchronous, in-memory implementation of the GraphStore protocol,
designed for building and querying context graphs from conversations and entities.

It formalizes context as a graph of connections, enabling meaningful connections between
concepts, entities, and conversations.

Key Features:
    - In-memory GraphStore implementation
    - Entity and relationship extraction from conversations
    - BFS-based neighbor discovery
    - Type-based indexing
    - Export to dictionary format
"""

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

from ..utils.exceptions import ProcessingError
from ..utils.logging import get_logger
from ..utils.progress_tracker import get_progress_tracker
from ..utils.types import EntityDict, RelationshipDict
from .entity_linker import EntityLinker

@dataclass
class ContextNode:
    """Context graph node (Internal implementation)."""
    node_id: str
    node_type: str
    content: str
    metadata: Dict[str, Any] = field(default_factory=dict)
    properties: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary format."""
        props = self.properties.copy()
        props.update(self.metadata)
        props["content"] = self.content
        return {
            "id": self.node_id,
            "type": self.node_type,
            "properties": props
        }

@dataclass
class ContextEdge:
    """Context graph edge (Internal implementation)."""
    source_id: str
    target_id: str
    edge_type: str
    weight: float = 1.0
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary format."""
        return {
            "source_id": self.source_id,
            "target_id": self.target_id,
            "type": self.edge_type,
            "weight": self.weight,
            "properties": self.metadata
        }

class ContextGraph:
    """
    In-memory implementation of context graph.
    
    Provides capabilities to build, store, and query a context graph.
    """

    def __init__(self, config: Optional[Dict[str, Any]] = None, **kwargs):
        """
        Initialize context graph.

        Args:
            config: Configuration dictionary
            **kwargs: Additional configuration options:
                - extract_entities: Extract entities from content (default: True)
                - extract_relationships: Extract relationships (default: True)
                - entity_linker: Entity linker instance
        """
        self.logger = get_logger("context_graph")
        self.config = config or {}
        self.config.update(kwargs)

        self.extract_entities = self.config.get("extract_entities", True)
        self.extract_relationships = self.config.get("extract_relationships", True)
        
        self.entity_linker = self.config.get("entity_linker") or EntityLinker()

        # Graph structure
        self.nodes: Dict[str, ContextNode] = {}
        self.edges: List[ContextEdge] = []
        
        # Adjacency list for efficient traversal: source_id -> list of edges
        self._adjacency: Dict[str, List[ContextEdge]] = defaultdict(list)

        # Indexes
        self.node_type_index: Dict[str, Set[str]] = defaultdict(set)
        self.edge_type_index: Dict[str, List[ContextEdge]] = defaultdict(list)
        
        # Progress tracker
        self.progress_tracker = get_progress_tracker()

    # --- GraphStore Protocol Implementation ---

    def add_nodes(self, nodes: List[Dict[str, Any]]) -> int:
        """
        Add nodes to graph.
        
        Args:
            nodes: List of nodes to add (dicts with id, type, properties)
            
        Returns:
            Number of nodes added
        """
        count = 0
        for node in nodes:
            # Extract content from properties if not explicit
            node_props = node.get("properties", {})
            content = node_props.get("content", node.get("id"))
            metadata = {k: v for k, v in node_props.items() if k != "content"}
            
            internal_node = ContextNode(
                node_id=node.get("id"),
                node_type=node.get("type", "entity"),
                content=content,
                metadata=metadata,
                properties=node_props
            )
            
            if self._add_internal_node(internal_node):
                count += 1
        return count

    def add_edges(self, edges: List[Dict[str, Any]]) -> int:
        """
        Add edges to graph.
        
        Args:
            edges: List of edges to add (dicts with source_id, target_id, type, weight, properties)
            
        Returns:
            Number of edges added
        """
        count = 0
        for edge in edges:
            internal_edge = ContextEdge(
                source_id=edge.get("source_id"),
                target_id=edge.get("target_id"),
                edge_type=edge.get("type", "related_to"),
                weight=edge.get("weight", 1.0),
                metadata=edge.get("properties", {})
            )
            
            if self._add_internal_edge(internal_edge):
                count += 1
        return count

    def get_neighbors(self, node_id: str, hops: int = 1) -> List[Dict[str, Any]]:
        """
        Get neighbors of a node.
        
        Returns list of dicts with neighbor info.
        """
        if node_id not in self.nodes:
            return []

        neighbors = []
        visited = {node_id}
        queue = deque([(node_id, 0)])  # (current_id, current_hop)

        while queue:
            current_id, current_hop = queue.popleft()
            
            if current_hop >= hops:
                continue

            # Get outgoing edges
            outgoing_edges = self._adjacency.get(current_id, [])
            for edge in outgoing_edges:
                neighbor_id = edge.target_id
                if neighbor_id not in visited:
                    visited.add(neighbor_id)
                    queue.append((neighbor_id, current_hop + 1))
                    
                    if neighbor_id in self.nodes:
                        node = self.nodes[neighbor_id]
                        neighbors.append({
                            "id": node.node_id,
                            "type": node.node_type,
                            "content": node.content,
                            "relationship": edge.edge_type,
                            "weight": edge.weight,
                            "hop": current_hop + 1
                        })

        return neighbors

    def query(self, query: str) -> List[Dict[str, Any]]:
        """
        Execute a simple keyword search query on the graph nodes.
        
        Args:
            query: Keyword query string
            
        Returns:
            List of matching node dicts
        """
        results = []
        query_lower = query.lower().split()
        
        for node in self.nodes.values():
            content_lower = node.content.lower()
            if any(word in content_lower for word in query_lower):
                # Calculate simple score
                overlap = sum(1 for word in query_lower if word in content_lower)
                score = overlap / len(query_lower) if query_lower else 0.0
                
                results.append({
                    "node": node.to_dict(),
                    "score": score,
                    "content": node.content
                })
        
        return sorted(results, key=lambda x: x["score"], reverse=True)

    def add_node(self, node_id: str, node_type: str, content: Optional[str] = None, **properties) -> bool:
        """
        Add a single node to the graph.
        
        Args:
            node_id: Unique identifier
            node_type: Node type (e.g., 'entity', 'concept')
            content: Node content/label
            **properties: Additional properties
        """
        content = content or node_id
        return self._add_internal_node(ContextNode(
            node_id=node_id,
            node_type=node_type,
            content=content,
            metadata=properties,
            properties=properties
        ))

    def add_edge(self, source_id: str, target_id: str, edge_type: str = "related_to", weight: float = 1.0, **properties) -> bool:
        """
        Add a single edge to the graph.
        
        Args:
            source_id: Source node ID
            target_id: Target node ID
            edge_type: Relationship type
            weight: Edge weight
            **properties: Additional properties
        """
        return self._add_internal_edge(ContextEdge(
            source_id=source_id,
            target_id=target_id,
            edge_type=edge_type,
            weight=weight,
            metadata=properties
        ))

    def find_node(self, node_id: str) -> Optional[Dict[str, Any]]:
        """Find a node by ID."""
        node = self.nodes.get(node_id)
        if node:
            return {
                "id": node.node_id,
                "type": node.node_type,
                "content": node.content,
                "metadata": node.metadata
            }
        return None

    def find_nodes(self, node_type: Optional[str] = None) -> List[Dict[str, Any]]:
        """Find nodes, optionally filtered by type."""
        if node_type:
            node_ids = self.node_type_index.get(node_type, set())
            nodes = [self.nodes[nid] for nid in node_ids]
        else:
            nodes = self.nodes.values()
            
        return [
            {
                "id": n.node_id,
                "type": n.node_type,
                "content": n.content,
                "metadata": n.metadata
            }
            for n in nodes
        ]

    def find_edges(self, edge_type: Optional[str] = None) -> List[Dict[str, Any]]:
        """Find edges, optionally filtered by type."""
        if edge_type:
            edges = self.edge_type_index.get(edge_type, [])
        else:
            edges = self.edges
            
        return [
            {
                "source": e.source_id,
                "target": e.target_id,
                "type": e.edge_type,
                "weight": e.weight,
                "metadata": e.metadata
            }
            for e in edges
        ]

    def stats(self) -> Dict[str, Any]:
        """Get graph statistics."""
        return {
            "node_count": len(self.nodes),
            "edge_count": len(self.edges),
            "node_types": {k: len(v) for k, v in self.node_type_index.items()},
            "edge_types": {k: len(v) for k, v in self.edge_type_index.items()},
            "density": self.density()
        }

    def density(self) -> float:
        """Calculate graph density."""
        n = len(self.nodes)
        if n < 2:
            return 0.0
        max_edges = n * (n - 1)  # Directed graph
        return len(self.edges) / max_edges

    # --- Internal Helpers ---

    def _add_internal_node(self, node: ContextNode) -> bool:
        """Internal method to add a node."""
        self.nodes[node.node_id] = node
        self.node_type_index[node.node_type].add(node.node_id)
        return True

    def _add_internal_edge(self, edge: ContextEdge) -> bool:
        """Internal method to add an edge."""
        # Ensure nodes exist
        if edge.source_id not in self.nodes:
            self._add_internal_node(ContextNode(edge.source_id, "entity", edge.source_id))
        if edge.target_id not in self.nodes:
            self._add_internal_node(ContextNode(edge.target_id, "entity", edge.target_id))
            
        self.edges.append(edge)
        self.edge_type_index[edge.edge_type].append(edge)
        self._adjacency[edge.source_id].append(edge)
        return True

    # --- Builder Methods (Legacy/Utility) ---

    def build_from_conversations(
        self,
        conversations: List[Union[str, Dict[str, Any]]],
        link_entities: bool = True,
        extract_intents: bool = False,
        extract_sentiments: bool = False,
        **options,
    ) -> Dict[str, Any]:
        """
        Build context graph from conversations and return dict representation.
        
        Args:
            conversations: List of conversation files or dictionaries
            ...
            
        Returns:
            Graph dictionary (nodes, edges)
        """
        tracking_id = self.progress_tracker.start_tracking(
            file=None,
            module="context",
            submodule="ContextGraph",
            message=f"Building graph from {len(conversations)} conversations",
        )

        try:
            for conv in conversations:
                conv_data = conv if isinstance(conv, dict) else self._load_conversation(conv)
                self._process_conversation(
                    conv_data, 
                    extract_intents=extract_intents,
                    extract_sentiments=extract_sentiments
                )

            if link_entities:
                self._link_entities()

            self.progress_tracker.stop_tracking(tracking_id, status="completed")
            return self.to_dict()

        except Exception as e:
            self.progress_tracker.stop_tracking(tracking_id, status="failed", message=str(e))
            raise

    def build_from_entities_and_relationships(
        self,
        entities: List[Dict[str, Any]],
        relationships: List[Dict[str, Any]],
        **kwargs
    ) -> Dict[str, Any]:
        """
        Build graph from entities and relationships.
        
        Args:
            entities: List of entity dictionaries
            relationships: List of relationship dictionaries
            **kwargs: Additional options
            
        Returns:
            Graph dictionary (nodes, edges)
        """
        tracking_id = self.progress_tracker.start_tracking(
            file=None,
            module="context",
            submodule="ContextGraph",
            message=f"Building graph from {len(entities)} entities and {len(relationships)} relationships",
        )

        try:
            # Add entities
            for entity in entities:
                entity_id = entity.get("id") or entity.get("entity_id")
                if entity_id:
                    self._add_internal_node(ContextNode(
                        node_id=entity_id,
                        node_type=entity.get("type", "entity"),
                        content=entity.get("text") or entity.get("label") or entity_id,
                        metadata=entity,
                        properties=entity
                    ))

            # Add relationships
            for rel in relationships:
                source = rel.get("source_id")
                target = rel.get("target_id")
                if source and target:
                    self._add_internal_edge(ContextEdge(
                        source_id=source,
                        target_id=target,
                        edge_type=rel.get("type", "related_to"),
                        weight=rel.get("confidence", 1.0),
                        metadata=rel
                    ))
            
            self.progress_tracker.stop_tracking(tracking_id, status="completed")
            return self.to_dict()

        except Exception as e:
            self.progress_tracker.stop_tracking(tracking_id, status="failed", message=str(e))
            raise

    def _process_conversation(self, conv_data: Dict[str, Any], **kwargs) -> None:
        """Process a single conversation."""
        conv_id = conv_data.get("id") or f"conv_{hash(str(conv_data)) % 10000}"
        
        # Add conversation node
        self._add_internal_node(ContextNode(
            node_id=conv_id,
            node_type="conversation",
            content=conv_data.get("content", "") or conv_data.get("summary", ""),
            metadata={"timestamp": conv_data.get("timestamp")}
        ))

        # Track name to ID mapping for relationship resolution
        name_to_id = {}

        # Extract entities
        if self.extract_entities:
            for entity in conv_data.get("entities", []):
                entity_id = entity.get("id") or entity.get("entity_id")
                entity_text = entity.get("text") or entity.get("label") or entity.get("name") or entity_id
                entity_type = entity.get("type", "entity")
                
                # Generate ID if missing
                if not entity_id and entity_text and self.entity_linker:
                    # Use EntityLinker to generate ID
                    if hasattr(self.entity_linker, "_generate_entity_id"):
                        entity_id = self.entity_linker._generate_entity_id(entity_text, entity_type)
                    else:
                        # Fallback ID generation
                        import hashlib
                        entity_hash = hashlib.md5(f"{entity_text}_{entity_type}".encode()).hexdigest()[:12]
                        entity_id = f"{entity_type.lower()}_{entity_hash}"
                
                if entity_id:
                    if entity_text:
                        name_to_id[entity_text] = entity_id
                        
                    self._add_internal_node(ContextNode(
                        node_id=entity_id,
                        node_type="entity",
                        content=entity_text,
                        metadata={"type": entity_type, **entity}
                    ))
                    self._add_internal_edge(ContextEdge(
                        source_id=conv_id,
                        target_id=entity_id,
                        edge_type="mentions"
                    ))

        # Extract relationships
        if self.extract_relationships:
            for rel in conv_data.get("relationships", []):
                source = rel.get("source_id")
                target = rel.get("target_id")
                
                # Resolve IDs from names if missing
                if not source and rel.get("source") and rel.get("source") in name_to_id:
                    source = name_to_id[rel.get("source")]
                
                if not target and rel.get("target") and rel.get("target") in name_to_id:
                    target = name_to_id[rel.get("target")]
                
                if source and target:
                    self._add_internal_edge(ContextEdge(
                        source_id=source,
                        target_id=target,
                        edge_type=rel.get("type", "related_to"),
                        weight=rel.get("confidence", 1.0)
                    ))

    def _link_entities(self) -> None:
        """Link similar entities using EntityLinker."""
        if not self.entity_linker:
            return
            
        entity_nodes = [n for n in self.nodes.values() if n.node_type == "entity"]
        for i, node1 in enumerate(entity_nodes):
            for node2 in entity_nodes[i+1:]:
                similarity = self.entity_linker._calculate_text_similarity(
                    node1.content.lower(), node2.content.lower()
                )
                if similarity >= self.entity_linker.similarity_threshold:
                    self._add_internal_edge(ContextEdge(
                        source_id=node1.node_id,
                        target_id=node2.node_id,
                        edge_type="similar_to",
                        weight=similarity
                    ))

    def _load_conversation(self, file_path: str) -> Dict[str, Any]:
        """Load conversation from file."""
        from ..utils.helpers import read_json_file
        from pathlib import Path
        return read_json_file(Path(file_path))

    def to_dict(self) -> Dict[str, Any]:
        """Export graph to dictionary format."""
        return {
            "nodes": [
                {
                    "id": n.node_id,
                    "type": n.node_type,
                    "content": n.content,
                    "metadata": n.metadata
                }
                for n in self.nodes.values()
            ],
            "edges": [
                {
                    "source": e.source_id,
                    "target": e.target_id,
                    "type": e.edge_type,
                    "weight": e.weight
                }
                for e in self.edges
            ],
            "statistics": {
                "node_count": len(self.nodes),
                "edge_count": len(self.edges)
            }
        }

# For backward compatibility
ContextGraphBuilder = ContextGraph
