"""
Entity Linker for Context Engineering

This module provides comprehensive entity linking capabilities for context
engineering, linking entities across different sources to build the web of
context. It assigns each entity a unique URL/URI and connects them meaningfully
to enable semantic understanding.

Algorithms Used:

URI Generation:
    - Hash-based URI: MD5 hash-based URI generation for entities without text
    - Text-based URI: URL-safe text encoding for entity text
    - Type-based URI: Entity type inclusion in URI
    - URI Registry: Dictionary-based URI-to-entity mapping

Entity Linking:
    - Text Similarity: Word overlap-based similarity calculation (Jaccard-like)
    - Knowledge Graph Lookup: Entity matching in knowledge graph
    - Similarity Threshold: Threshold-based entity matching
    - Cross-Document Linking: Entity linking across multiple documents
    - Bidirectional Linking: Symmetric relationship creation

Entity Web Construction:
    - Graph Building: Graph construction from entity links
    - Link Aggregation: Link counting and statistics
    - Entity Registry: Entity-to-URI mapping
    - Link Storage: Dictionary-based link storage (entity_id -> List[EntityLink])

Key Features:
    - Links entities across different sources
    - Assigns unique identifiers (URLs/URIs) to entities
    - Creates semantic connections between entities
    - Builds entity connection web
    - Supports cross-document entity linking
    - Enables entity disambiguation and resolution
    - Similarity-based entity matching
    - Bidirectional entity linking
    - Entity web construction and statistics
    - Configurable similarity thresholds

Main Classes:
    - EntityLink: Entity link data structure with source_entity_id, target_entity_id, link_type, confidence, source, metadata
    - LinkedEntity: Linked entity with context including entity_id, uri, text, type, linked_entities, context, confidence
    - EntityLinker: Entity linker for context engineering

Example Usage:
    >>> from semantica.context import EntityLinker
    >>> linker = EntityLinker(knowledge_graph=kg)
    >>> uri = linker.assign_uri("entity_1", "Python", "PROGRAMMING_LANGUAGE")
    >>> linked_entities = linker.link("Python is a programming language", entities=entities)
    >>> linker.link_entities("entity_1", "entity_2", "related_to", confidence=0.9)
    >>> web = linker.build_entity_web()

Author: Semantica Contributors
License: MIT
"""

import hashlib
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Set, Tuple, Union
from urllib.parse import quote

from ..utils.exceptions import ProcessingError, ValidationError
from ..utils.logging import get_logger
from ..utils.progress_tracker import get_progress_tracker
from ..utils.types import EntityDict


@dataclass
class EntityLink:
    """Entity link definition."""

    source_entity_id: str
    target_entity_id: str
    link_type: str  # "same_as", "related_to", "part_of", etc.
    confidence: float = 1.0
    source: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class LinkedEntity:
    """Linked entity with context."""

    entity_id: str
    uri: str
    text: str
    type: str
    linked_entities: List[EntityLink] = field(default_factory=list)
    context: Dict[str, Any] = field(default_factory=dict)
    confidence: float = 1.0


class EntityLinker:
    """
    Entity linker for context engineering.

    • Links entities across different sources
    • Assigns unique identifiers (URLs/URIs) to entities
    • Creates semantic connections between entities
    • Builds entity connection web
    • Supports cross-document entity linking
    • Enables entity disambiguation and resolution
    """

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

        Args:
            config: Configuration dictionary
            **kwargs: Additional configuration options:
                - knowledge_graph: Knowledge graph for entity lookup
                - similarity_threshold: Similarity threshold for linking (default: 0.8)
                - base_uri: Base URI for entity URIs
                - enable_cross_document_linking: Enable cross-document linking (default: True)
        """
        self.logger = get_logger("entity_linker")
        self.config = config or {}
        self.config.update(kwargs)

        self.knowledge_graph = self.config.get("knowledge_graph")
        self.similarity_threshold = self.config.get("similarity_threshold", 0.8)
        self.base_uri = self.config.get("base_uri", "https://semantica.dev/entity/")
        self.enable_cross_document_linking = self.config.get(
            "enable_cross_document_linking", True
        )

        # Initialize progress tracker
        self.progress_tracker = get_progress_tracker()

        # Entity registry: entity_id -> URI
        self.entity_registry: Dict[str, str] = {}

        # Entity links: entity_id -> List[EntityLink]
        self.entity_links: Dict[str, List[EntityLink]] = {}

    def assign_uri(
        self,
        entity_id: str,
        entity_text: Optional[str] = None,
        entity_type: Optional[str] = None,
    ) -> str:
        """
        Assign unique URI to entity.

        Args:
            entity_id: Entity identifier
            entity_text: Entity text/name
            entity_type: Entity type

        Returns:
            Entity URI
        """
        if entity_id in self.entity_registry:
            return self.entity_registry[entity_id]

        # Generate URI from entity information
        if entity_text:
            # Create URI-safe identifier
            uri_safe = quote(entity_text.lower().replace(" ", "_"), safe="")
            uri = f"{self.base_uri}{uri_safe}"
        else:
            # Use hash of entity_id
            entity_hash = hashlib.md5(entity_id.encode()).hexdigest()[:8]
            uri = f"{self.base_uri}{entity_hash}"

        # Add type if available
        if entity_type:
            uri = f"{uri}#{entity_type.lower()}"

        self.entity_registry[entity_id] = uri
        self.logger.debug(f"Assigned URI {uri} to entity {entity_id}")

        return uri

    def link(
        self,
        text: str,
        entities: Optional[List[EntityDict]] = None,
        context: Optional[List[Dict[str, Any]]] = None,
    ) -> List[LinkedEntity]:
        """
        Link entities in text to knowledge graph.

        Args:
            text: Input text
            entities: List of entities to link
            context: Additional context information

        Returns:
            List of linked entities
        """
        # Track entity linking
        tracking_id = self.progress_tracker.start_tracking(
            file=None,
            module="context",
            submodule="EntityLinker",
            message=f"Linking entities in text",
        )

        try:
            if not entities:
                entities = []

            linked_entities = []

            self.progress_tracker.update_tracking(
                tracking_id, message=f"Linking {len(entities)} entities..."
            )

            for entity in entities:
                entity_id = entity.get("id") or entity.get("entity_id")
                entity_text = (
                    entity.get("text") or entity.get("label") or entity.get("name", "")
                )
                entity_type = entity.get("type") or entity.get("entity_type", "UNKNOWN")

                if not entity_id:
                    # Generate entity_id if not present
                    entity_id = self._generate_entity_id(entity_text, entity_type)

                # Assign URI
                uri = self.assign_uri(entity_id, entity_text, entity_type)

                # Find linked entities
                linked_entity_links = self._find_linked_entities(
                    entity_id, entity_text, entity_type, entities, context
                )

                linked_entity = LinkedEntity(
                    entity_id=entity_id,
                    uri=uri,
                    text=entity_text,
                    type=entity_type,
                    linked_entities=linked_entity_links,
                    context=entity.get("metadata", {}),
                    confidence=entity.get("confidence", 1.0),
                )

                linked_entities.append(linked_entity)

                # Store links
                if entity_id not in self.entity_links:
                    self.entity_links[entity_id] = []
                self.entity_links[entity_id].extend(linked_entity_links)

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

        return linked_entities

    def link_entities(
        self,
        entity1_id: str,
        entity2_id: str,
        link_type: str = "related_to",
        confidence: float = 1.0,
        source: Optional[str] = None,
        **metadata,
    ) -> bool:
        """
        Create explicit link between two entities.

        Args:
            entity1_id: First entity ID
            entity2_id: Second entity ID
            link_type: Link type
            confidence: Link confidence
            source: Source of the link
            **metadata: Additional metadata

        Returns:
            True if link created successfully
        """
        # Track entity linking
        tracking_id = self.progress_tracker.start_tracking(
            file=None,
            module="context",
            submodule="EntityLinker",
            message=f"Linking {entity1_id} -> {entity2_id}",
        )

        try:
            link = EntityLink(
                source_entity_id=entity1_id,
                target_entity_id=entity2_id,
                link_type=link_type,
                confidence=confidence,
                source=source,
                metadata=metadata,
            )

            if entity1_id not in self.entity_links:
                self.entity_links[entity1_id] = []

            self.entity_links[entity1_id].append(link)

            # Create bidirectional link if appropriate
            if link_type in ["related_to", "same_as"]:
                reverse_link = EntityLink(
                    source_entity_id=entity2_id,
                    target_entity_id=entity1_id,
                    link_type=link_type,
                    confidence=confidence,
                    source=source,
                    metadata=metadata,
                )

                if entity2_id not in self.entity_links:
                    self.entity_links[entity2_id] = []
                self.entity_links[entity2_id].append(reverse_link)

            self.logger.debug(f"Linked {entity1_id} --{link_type}--> {entity2_id}")

            self.progress_tracker.stop_tracking(
                tracking_id,
                status="completed",
                message=f"Linked {entity1_id} -> {entity2_id}",
            )
            return True

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

    def get_entity_links(self, entity_id: str) -> List[EntityLink]:
        """
        Get all links for an entity.

        Args:
            entity_id: Entity identifier

        Returns:
            List of entity links
        """
        return self.entity_links.get(entity_id, [])

    def get_entity_uri(self, entity_id: str) -> Optional[str]:
        """
        Get URI for an entity.

        Args:
            entity_id: Entity identifier

        Returns:
            Entity URI or None if not found
        """
        return self.entity_registry.get(entity_id)

    def find_similar_entities(
        self,
        entity_text: str,
        entity_type: Optional[str] = None,
        threshold: Optional[float] = None,
    ) -> List[Tuple[str, float]]:
        """
        Find similar entities in knowledge graph.

        Args:
            entity_text: Entity text to search for
            entity_type: Optional entity type filter
            threshold: Similarity threshold (uses default if None)

        Returns:
            List of (entity_id, similarity_score) tuples
        """
        threshold = threshold or self.similarity_threshold

        if not self.knowledge_graph:
            return []

        # Simple text-based similarity (can be enhanced with embeddings)
        similar_entities = []

        entities = self.knowledge_graph.get("entities", [])

        for entity in entities:
            if entity_type and entity.get("type") != entity_type:
                continue

            entity_text2 = (
                entity.get("text") or entity.get("label") or entity.get("name", "")
            )

            # Calculate simple similarity (can be enhanced)
            similarity = self._calculate_text_similarity(
                entity_text.lower(), entity_text2.lower()
            )

            if similarity >= threshold:
                entity_id = entity.get("id") or entity.get("entity_id")
                if entity_id:
                    similar_entities.append((entity_id, similarity))

        # Sort by similarity
        similar_entities.sort(key=lambda x: x[1], reverse=True)

        return similar_entities

    def _find_linked_entities(
        self,
        entity_id: str,
        entity_text: str,
        entity_type: str,
        all_entities: List[EntityDict],
        context: Optional[List[Dict[str, Any]]],
    ) -> List[EntityLink]:
        """Find entities linked to this entity."""
        links = []

        # Find similar entities in knowledge graph
        if self.knowledge_graph:
            similar = self.find_similar_entities(entity_text, entity_type)
            for similar_id, similarity in similar:
                if similar_id != entity_id:
                    links.append(
                        EntityLink(
                            source_entity_id=entity_id,
                            target_entity_id=similar_id,
                            link_type="same_as" if similarity >= 0.9 else "related_to",
                            confidence=similarity,
                            metadata={"similarity": similarity},
                        )
                    )

        # Find relationships in context
        if context:
            for ctx in context:
                relationships = ctx.get("relationships", [])
                for rel in relationships:
                    source_id = rel.get("source_id")
                    target_id = rel.get("target_id")

                    if source_id == entity_id:
                        links.append(
                            EntityLink(
                                source_entity_id=entity_id,
                                target_entity_id=target_id,
                                link_type=rel.get("type", "related_to"),
                                confidence=rel.get("confidence", 1.0),
                                source=ctx.get("source"),
                            )
                        )
                    elif target_id == entity_id:
                        links.append(
                            EntityLink(
                                source_entity_id=entity_id,
                                target_entity_id=source_id,
                                link_type=rel.get("type", "related_to"),
                                confidence=rel.get("confidence", 1.0),
                                source=ctx.get("source"),
                            )
                        )

        return links

    def _calculate_text_similarity(self, text1: str, text2: str) -> float:
        """Calculate simple text similarity."""
        if not text1 or not text2:
            return 0.0

        # Simple word overlap similarity
        words1 = set(text1.split())
        words2 = set(text2.split())

        if not words1 or not words2:
            return 0.0

        intersection = words1 & words2
        union = words1 | words2

        if not union:
            return 0.0

        return len(intersection) / len(union)

    def _generate_entity_id(self, text: str, entity_type: str) -> str:
        """Generate entity ID from text and type."""
        entity_hash = hashlib.md5(f"{text}_{entity_type}".encode()).hexdigest()[:12]
        return f"{entity_type.lower()}_{entity_hash}"

    def build_entity_web(self) -> Dict[str, Any]:
        """
        Build entity connection web.

        Returns:
            Entity web dictionary
        """
        web = {
            "entities": {},
            "links": [],
            "statistics": {
                "total_entities": len(self.entity_registry),
                "total_links": sum(len(links) for links in self.entity_links.values()),
            },
        }

        # Add entities
        for entity_id, uri in self.entity_registry.items():
            web["entities"][entity_id] = {
                "uri": uri,
                "links": len(self.entity_links.get(entity_id, [])),
            }

        # Add links
        for entity_id, links in self.entity_links.items():
            for link in links:
                web["links"].append(
                    {
                        "source": link.source_entity_id,
                        "target": link.target_entity_id,
                        "type": link.link_type,
                        "confidence": link.confidence,
                    }
                )

        return web

    # Linking Methods
    def link_text(self, text: str, entities: Optional[List[EntityDict]] = None) -> List[LinkedEntity]:
        """
        Link entities in text.
        
        Args:
            text: Text containing entities
            entities: List of entities to link (optional)
        
        Returns:
            List of LinkedEntity objects
        
        Example:
            >>> linked = linker.link_text("Python is used for ML", entities=[...])
        """
        return self.link(text, entities=entities or [])

    def link_batch(self, entities: List[EntityDict]) -> List[LinkedEntity]:
        """
        Link multiple entities.
        
        Args:
            entities: List of entities to link
        
        Returns:
            List of LinkedEntity objects
        
        Example:
            >>> linked = linker.link_batch(entities)
        """
        linked_results = []
        for entity in entities:
            uri = self.assign_uri(
                entity.get("id", ""),
                entity.get("text"),
                entity.get("type")
            )
            
            # Find similar entities
            similar = self.find_similar_entities(
                entity.get("text", ""),
                entity.get("type", ""),
                threshold=self.similarity_threshold
            )
            
            linked_entities = []
            for similar_id, similarity in similar:
                linked_entities.append(
                    EntityLink(
                        source_entity_id=entity.get("id", ""),
                        target_entity_id=similar_id,
                        link_type="similar_to",
                        confidence=similarity
                    )
                )
            
            linked_results.append(
                LinkedEntity(
                    entity_id=entity.get("id", ""),
                    uri=uri,
                    text=entity.get("text", ""),
                    type=entity.get("type", ""),
                    linked_entities=linked_entities,
                    confidence=1.0
                )
            )
        
        return linked_results

    # Search Methods
    def find_similar(self, entity: Union[str, EntityDict], threshold: float = 0.8) -> List[Tuple[str, float]]:
        """
        Find similar entities.
        
        Args:
            entity: Entity text or EntityDict
            threshold: Similarity threshold (default: 0.8)
        
        Returns:
            List of (entity_id, similarity) tuples
        
        Example:
            >>> similar = linker.find_similar("Python", threshold=0.8)
        """
        if isinstance(entity, dict):
            text = entity.get("text", "")
            entity_type = entity.get("type", "")
        else:
            text = entity
            entity_type = ""
        
        return self.find_similar_entities(text, entity_type, threshold=threshold)

    def find_by_text(self, text: str, threshold: float = 0.8) -> List[Dict[str, Any]]:
        """
        Find by text.
        
        Args:
            text: Text to search for
            threshold: Similarity threshold (default: 0.8)
        
        Returns:
            List of matching entity dicts
        
        Example:
            >>> entities = linker.find_by_text("Python", threshold=0.8)
        """
        results = []
        for entity_id, uri in self.entity_registry.items():
            # Get entity text from knowledge graph or registry
            if self.knowledge_graph:
                nodes = self.knowledge_graph.get("nodes", [])
                for node in nodes:
                    if node.get("id") == entity_id:
                        similarity = self._calculate_similarity(text, node.get("content", ""))
                        if similarity >= threshold:
                            results.append({
                                "entity_id": entity_id,
                                "uri": uri,
                                "text": node.get("content", ""),
                                "type": node.get("type", ""),
                                "similarity": similarity,
                            })
                        break
        return sorted(results, key=lambda x: x.get("similarity", 0), reverse=True)

    def find_by_type(self, type: str, limit: int = 10) -> List[Dict[str, Any]]:
        """
        Find by type.
        
        Args:
            type: Entity type
            limit: Maximum results (default: 10)
        
        Returns:
            List of entity dicts of specified type
        
        Example:
            >>> entities = linker.find_by_type("PROGRAMMING_LANGUAGE", limit=10)
        """
        results = []
        if self.knowledge_graph:
            nodes = self.knowledge_graph.get("nodes", [])
            for node in nodes:
                if node.get("type") == type:
                    entity_id = node.get("id", "")
                    uri = self.entity_registry.get(entity_id, "")
                    results.append({
                        "entity_id": entity_id,
                        "uri": uri,
                        "text": node.get("content", ""),
                        "type": type,
                    })
                    if len(results) >= limit:
                        break
        return results

    def find_related(self, entity_id: str, max_hops: int = 2) -> List[Dict[str, Any]]:
        """
        Find related entities.
        
        Args:
            entity_id: Entity ID
            max_hops: Maximum hops (default: 2)
        
        Returns:
            List of related entity dicts
        
        Example:
            >>> related = linker.find_related("entity_123", max_hops=2)
        """
        related = []
        visited = {entity_id}
        current_level = {entity_id}
        
        for hop in range(max_hops):
            next_level = set()
            for current_id in current_level:
                links = self.entity_links.get(current_id, [])
                for link in links:
                    target_id = link.target_entity_id
                    if target_id not in visited:
                        visited.add(target_id)
                        next_level.add(target_id)
                        
                        # Get entity info
                        uri = self.entity_registry.get(target_id, "")
                        if self.knowledge_graph:
                            nodes = self.knowledge_graph.get("nodes", [])
                            for node in nodes:
                                if node.get("id") == target_id:
                                    related.append({
                                        "entity_id": target_id,
                                        "uri": uri,
                                        "text": node.get("content", ""),
                                        "type": node.get("type", ""),
                                        "link_type": link.link_type,
                                        "confidence": link.confidence,
                                        "hop": hop + 1,
                                    })
                                    break
            
            current_level = next_level
        
        return related

    # URI Methods
    def get_uri(self, entity_id: str) -> Optional[str]:
        """
        Get URI for entity.
        
        Args:
            entity_id: Entity ID
        
        Returns:
            URI or None if not found
        
        Example:
            >>> uri = linker.get_uri("entity_123")
        """
        return self.entity_registry.get(entity_id)

    def get_all_uris(self) -> Dict[str, str]:
        """
        Get all URIs.
        
        Returns:
            Dict mapping entity_id to URI
        
        Example:
            >>> uris = linker.get_all_uris()
        """
        return self.entity_registry.copy()

    def resolve_uri(self, uri: str) -> Optional[Dict[str, Any]]:
        """
        Resolve URI to entity.
        
        Args:
            uri: URI to resolve
        
        Returns:
            Entity dict or None if not found
        
        Example:
            >>> entity = linker.resolve_uri("https://semantica.dev/entity/python")
        """
        # Find entity_id by URI
        for entity_id, entity_uri in self.entity_registry.items():
            if entity_uri == uri:
                # Get entity info
                if self.knowledge_graph:
                    nodes = self.knowledge_graph.get("nodes", [])
                    for node in nodes:
                        if node.get("id") == entity_id:
                            return {
                                "entity_id": entity_id,
                                "uri": uri,
                                "text": node.get("content", ""),
                                "type": node.get("type", ""),
                            }
                return {
                    "entity_id": entity_id,
                    "uri": uri,
                }
        return None

    # Statistics
    def link_count(self, entity_id: str) -> int:
        """
        Get link count for entity.
        
        Args:
            entity_id: Entity ID
        
        Returns:
            Number of links for entity
        
        Example:
            >>> count = linker.link_count("entity_123")
        """
        return len(self.entity_links.get(entity_id, []))

    def stats(self) -> Dict[str, Any]:
        """
        Get statistics.
        
        Returns:
            Statistics dict
        
        Example:
            >>> stats = linker.stats()
        """
        web = self.build_entity_web()
        return {
            "total_entities": web["statistics"]["total_entities"],
            "total_links": web["statistics"]["total_links"],
            "average_links_per_entity": (
                web["statistics"]["total_links"] / web["statistics"]["total_entities"]
                if web["statistics"]["total_entities"] > 0 else 0
            ),
        }

    def most_linked(self, limit: int = 10) -> List[Dict[str, Any]]:
        """
        Get most linked entities.
        
        Args:
            limit: Maximum results (default: 10)
        
        Returns:
            List of entity dicts sorted by link count
        
        Example:
            >>> top = linker.most_linked(limit=10)
        """
        entity_link_counts = []
        for entity_id, links in self.entity_links.items():
            uri = self.entity_registry.get(entity_id, "")
            entity_link_counts.append({
                "entity_id": entity_id,
                "uri": uri,
                "link_count": len(links),
            })
        
        # Sort by link count
        entity_link_counts.sort(key=lambda x: x["link_count"], reverse=True)
        
        # Add entity text if available
        if self.knowledge_graph:
            nodes = self.knowledge_graph.get("nodes", [])
            node_map = {node.get("id"): node for node in nodes}
            for item in entity_link_counts:
                node = node_map.get(item["entity_id"])
                if node:
                    item["text"] = node.get("content", "")
                    item["type"] = node.get("type", "")
        
        return entity_link_counts[:limit]