"""
Entity Normalizer - ensures consistent entity representation.

This module provides always-on normalization that:
1. Lowercases and trims all entities
2. Resolves common aliases (e.g., "US" -> "united states")
3. Deduplicates within documents
"""
from __future__ import annotations
import re
from dataclasses import dataclass
from typing import Optional


@dataclass
class NormalizedEntity:
    """A normalized entity with canonical form and aliases."""
    canonical: str  # The normalized, canonical form
    original: str   # The original form as found in text
    entity_type: str = "ENTITY"

    def __hash__(self):
        return hash(self.canonical)

    def __eq__(self, other):
        if isinstance(other, NormalizedEntity):
            return self.canonical == other.canonical
        return False


class EntityNormalizer:
    """
    Always-on entity normalizer that ensures consistency.

    Features:
    - Case normalization (lowercase canonical form)
    - Alias resolution (common abbreviations)
    - Whitespace cleanup
    - Duplicate detection
    """

    # Common aliases for normalization
    ALIASES: dict[str, str] = {
        # Countries
        "us": "united states",
        "usa": "united states",
        "u.s.": "united states",
        "u.s.a.": "united states",
        "uk": "united kingdom",
        "u.k.": "united kingdom",

        # Companies
        "fb": "facebook",
        "meta": "meta platforms",
        "goog": "google",
        "googl": "google",
        "msft": "microsoft",
        "amzn": "amazon",
        "aapl": "apple",

        # Common abbreviations
        "dr.": "doctor",
        "mr.": "mister",
        "mrs.": "missus",
        "ms.": "miss",
        "prof.": "professor",
        "jr.": "junior",
        "sr.": "senior",
    }

    def __init__(self, custom_aliases: Optional[dict[str, str]] = None):
        """
        Initialize normalizer with optional custom aliases.

        Args:
            custom_aliases: Additional alias mappings
        """
        self.aliases = {**self.ALIASES}
        if custom_aliases:
            self.aliases.update({
                k.lower(): v.lower() for k, v in custom_aliases.items()
            })

        # Cache for normalized forms
        self._cache: dict[str, str] = {}

    def normalize(self, entity: str) -> NormalizedEntity:
        """
        Normalize a single entity string.

        Args:
            entity: Raw entity string

        Returns:
            NormalizedEntity with canonical form
        """
        original = entity

        # Check cache
        if entity in self._cache:
            return NormalizedEntity(
                canonical=self._cache[entity],
                original=original
            )

        # Step 1: Lowercase and strip
        normalized = entity.lower().strip()

        # Step 2: Clean whitespace
        normalized = re.sub(r'\s+', ' ', normalized)

        # Step 3: Remove common suffixes that don't add meaning
        normalized = re.sub(r"'s$", "", normalized)  # Remove possessive

        # Step 4: Apply alias resolution
        if normalized in self.aliases:
            normalized = self.aliases[normalized]

        # Step 5: Handle multi-word entities
        words = normalized.split()
        resolved_words = []
        for word in words:
            if word in self.aliases:
                resolved_words.append(self.aliases[word])
            else:
                resolved_words.append(word)
        normalized = ' '.join(resolved_words)

        # Cache and return
        self._cache[entity] = normalized
        return NormalizedEntity(canonical=normalized, original=original)

    def normalize_many(self, entities: list[str]) -> list[NormalizedEntity]:
        """
        Normalize multiple entities, removing duplicates.

        Args:
            entities: List of raw entity strings

        Returns:
            Deduplicated list of normalized entities
        """
        seen: set[str] = set()
        result: list[NormalizedEntity] = []

        for entity in entities:
            normalized = self.normalize(entity)
            if normalized.canonical not in seen:
                seen.add(normalized.canonical)
                result.append(normalized)

        return result

    def normalize_relations(
        self,
        relations: list[tuple[str, str, str]]
    ) -> list[tuple[NormalizedEntity, str, NormalizedEntity]]:
        """
        Normalize all entities in relation triples.

        Args:
            relations: List of (subject, relation, object) tuples

        Returns:
            Normalized relation triples
        """
        result = []
        for subj, rel, obj in relations:
            norm_subj = self.normalize(subj)
            norm_obj = self.normalize(obj)
            # Relation types are also normalized
            norm_rel = rel.lower().strip().replace(' ', '_')
            result.append((norm_subj, norm_rel, norm_obj))
        return result

    def add_alias(self, alias: str, canonical: str) -> None:
        """
        Add a custom alias at runtime.

        Args:
            alias: The alias form
            canonical: The canonical form it should resolve to
        """
        self.aliases[alias.lower()] = canonical.lower()
        # Clear cache to ensure new alias is applied
        self._cache.clear()

    def get_canonical(self, entity: str) -> str:
        """
        Get just the canonical form of an entity.

        Args:
            entity: Raw entity string

        Returns:
            Canonical (normalized) form
        """
        return self.normalize(entity).canonical


# Global default normalizer instance
_default_normalizer: Optional[EntityNormalizer] = None


def get_normalizer() -> EntityNormalizer:
    """Get the global default normalizer instance."""
    global _default_normalizer
    if _default_normalizer is None:
        _default_normalizer = EntityNormalizer()
    return _default_normalizer
