"""
Entity Linker - cross-references entities across documents.

This module provides:
1. Global entity registry for deduplication
2. Cross-document entity linking
3. Entity disambiguation
"""
from __future__ import annotations
import threading
from dataclasses import dataclass, field
from typing import Optional
from collections import defaultdict

from .normalizer import EntityNormalizer, get_normalizer


@dataclass
class LinkedEntity:
    """An entity linked across multiple documents."""
    canonical: str
    aliases: set[str] = field(default_factory=set)
    episode_ids: set[str] = field(default_factory=set)
    mention_count: int = 0

    def add_mention(self, episode_id: str, original_form: str):
        """Record a new mention of this entity."""
        self.aliases.add(original_form)
        self.episode_ids.add(episode_id)
        self.mention_count += 1


class EntityLinker:
    """
    Cross-document entity linker.

    Maintains a global registry of entities and links them
    across documents automatically.
    """

    def __init__(self, normalizer: Optional[EntityNormalizer] = None):
        """
        Initialize the entity linker.

        Args:
            normalizer: Entity normalizer instance (uses global if not provided)
        """
        self.normalizer = normalizer or get_normalizer()

        # Canonical form -> LinkedEntity
        self._registry: dict[str, LinkedEntity] = {}

        # Episode ID -> set of canonical entities
        self._episode_entities: dict[str, set[str]] = defaultdict(set)

        # Thread safety
        self._lock = threading.RLock()

    def link(
        self,
        entities: list[str],
        episode_id: str
    ) -> list[LinkedEntity]:
        """
        Link entities from a document to the global registry.

        Args:
            entities: Raw entity strings
            episode_id: ID of the episode containing these entities

        Returns:
            List of LinkedEntity objects
        """
        result = []

        with self._lock:
            for entity in entities:
                normalized = self.normalizer.normalize(entity)
                canonical = normalized.canonical

                if canonical not in self._registry:
                    # New entity
                    self._registry[canonical] = LinkedEntity(canonical=canonical)

                linked = self._registry[canonical]
                linked.add_mention(episode_id, normalized.original)
                self._episode_entities[episode_id].add(canonical)
                result.append(linked)

        return result

    def get_linked_episodes(self, entity: str) -> set[str]:
        """
        Find all episodes containing an entity (or its aliases).

        Args:
            entity: Raw entity string

        Returns:
            Set of episode IDs
        """
        canonical = self.normalizer.get_canonical(entity)

        with self._lock:
            if canonical in self._registry:
                return self._registry[canonical].episode_ids.copy()
        return set()

    def get_related_entities(self, entity: str) -> list[str]:
        """
        Find entities that co-occur with the given entity.

        Args:
            entity: Raw entity string

        Returns:
            List of co-occurring canonical entity names
        """
        canonical = self.normalizer.get_canonical(entity)

        with self._lock:
            if canonical not in self._registry:
                return []

            # Get all episodes containing this entity
            episode_ids = self._registry[canonical].episode_ids

            # Find all other entities in those episodes
            related: set[str] = set()
            for ep_id in episode_ids:
                for other_entity in self._episode_entities.get(ep_id, set()):
                    if other_entity != canonical:
                        related.add(other_entity)

            return list(related)

    def get_entity(self, entity: str) -> Optional[LinkedEntity]:
        """
        Get linked entity info by name.

        Args:
            entity: Raw entity string

        Returns:
            LinkedEntity or None if not found
        """
        canonical = self.normalizer.get_canonical(entity)
        with self._lock:
            return self._registry.get(canonical)

    def get_all_entities(self) -> list[LinkedEntity]:
        """Get all registered entities."""
        with self._lock:
            return list(self._registry.values())

    def get_stats(self) -> dict:
        """Get linker statistics."""
        with self._lock:
            return {
                "total_entities": len(self._registry),
                "total_episodes": len(self._episode_entities),
                "avg_mentions_per_entity": (
                    sum(e.mention_count for e in self._registry.values()) /
                    max(len(self._registry), 1)
                ),
            }


# Global default linker instance
_default_linker: Optional[EntityLinker] = None


def get_linker() -> EntityLinker:
    """Get the global default entity linker instance."""
    global _default_linker
    if _default_linker is None:
        _default_linker = EntityLinker()
    return _default_linker
