"""
Coreference Resolution Module

This module provides comprehensive coreference resolution capabilities for resolving
pronoun references and entity coreferences in text, enabling better understanding
of entity mentions and references. Supports multiple extraction methods for
underlying entity extraction.

Supported Methods (for underlying NER extractors):
    - "pattern": Pattern-based extraction
    - "regex": Regex-based extraction
    - "rules": Rule-based extraction
    - "ml": ML-based extraction (spaCy)
    - "huggingface": HuggingFace model extraction
    - "llm": LLM-based extraction
    - Any method supported by NERExtractor

Algorithms Used:
    - Pronoun Resolution: Rule-based and distance-based antecedent resolution
    - Entity Coreference: String matching and similarity-based coreference detection
    - Coreference Chain Building: Graph-based chain construction algorithms
    - Mention Extraction: Pattern-based and ML-based mention detection
    - Ambiguity Resolution: Context-aware disambiguation algorithms
    - Distance Metrics: Character distance and sentence distance calculations

Key Features:
    - Pronoun resolution (he, she, it, they, etc.)
    - Entity coreference detection
    - Coreference chain construction
    - Ambiguity resolution
    - Cross-document coreference support
    - Mention extraction and tracking
    - Integration with multiple NER extraction methods
    - Method parameter support for underlying extractors

Main Classes:
    - CoreferenceResolver: Main coreference resolution coordinator
    - PronounResolver: Pronoun resolution engine
    - EntityCoreferenceDetector: Entity coreference detection
    - CoreferenceChainBuilder: Coreference chain construction
    - Mention: Mention representation dataclass
    - CoreferenceChain: Coreference chain representation dataclass

Example Usage:
    >>> from semantica.semantic_extract import CoreferenceResolver
    >>> # Using default methods
    >>> resolver = CoreferenceResolver()
    >>> chains = resolver.resolve_coreferences("John went to the store. He bought milk.")
    >>> 
    >>> # Using LLM-based extraction for entities
    >>> resolver = CoreferenceResolver(method="llm", provider="openai")
    >>> chains = resolver.resolve_coreferences("John went to the store. He bought milk.")
    >>> 
    >>> pronouns = resolver.resolve_pronouns("Mary said she would come.")

Author: Semantica Contributors
License: MIT
"""

from dataclasses import dataclass, field
import re
from typing import Any, Dict, List, Optional, Tuple, Union

from ..utils.exceptions import ProcessingError
from ..utils.logging import get_logger
from ..utils.progress_tracker import get_progress_tracker
from .ner_extractor import Entity


@dataclass
class Mention:
    """Mention representation."""
    
    text: str
    start_char: int
    end_char: int
    mention_type: str  # pronoun, entity, nominal
    entity_id: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class CoreferenceChain:
    """Coreference chain representation."""
    
    mentions: List[Mention]
    representative: Mention
    entity_type: Optional[str] = None


class CoreferenceResolver:
    """Coreference resolution handler."""
    
    def __init__(self, method: Union[str, List[str]] = None, config=None, **kwargs):
        """
        Initialize coreference resolver.
        
        Args:
            method: Extraction method(s) for underlying NER extractors.
                   Can be passed to ner_method in config.
            config: Legacy config dict (deprecated, use kwargs)
            **kwargs: Configuration options:
                - ner_method: Method for NER extraction (if entities need to be extracted)
                - Other options passed to sub-components
        """
        self.logger = get_logger("coreference_resolver")
        self.config = config or {}
        self.config.update(kwargs)
        self.progress_tracker = get_progress_tracker()
        
        # Store method for passing to extractors if needed
        if method is not None:
            self.config["ner_method"] = method
        
        self.pronoun_resolver = PronounResolver(**self.config.get("pronoun", {}))
        self.entity_detector = EntityCoreferenceDetector(**self.config.get("entity", {}))
        self.chain_builder = CoreferenceChainBuilder(**self.config.get("chain", {}))
    
    def resolve_coreferences(self, text: str, **options) -> List[CoreferenceChain]:
        """
        Resolve coreferences in text.
        
        Args:
            text: Input text
            **options: Resolution options
            
        Returns:
            list: List of coreference chains
        """
        tracking_id = self.progress_tracker.start_tracking(
            module="semantic_extract",
            submodule="CoreferenceResolver",
            message="Resolving coreferences in text"
        )
        
        try:
            # Extract mentions
            self.progress_tracker.update_tracking(tracking_id, message="Extracting mentions...")
            mentions = self._extract_mentions(text)
            
            # Resolve pronouns
            self.progress_tracker.update_tracking(tracking_id, message="Resolving pronouns...")
            pronoun_resolutions = self.pronoun_resolver.resolve_pronouns(text, mentions, **options)
            
            # Detect entity coreferences
            self.progress_tracker.update_tracking(tracking_id, message="Detecting entity coreferences...")
            entity_corefs = self.entity_detector.detect_entity_coreferences(text, mentions, **options)
            
            # Build chains
            self.progress_tracker.update_tracking(tracking_id, message="Building coreference chains...")
            chains = self.chain_builder.build_coreference_chains(mentions, **options)
            
            self.progress_tracker.stop_tracking(tracking_id, status="completed",
                                              message=f"Resolved {len(chains)} coreference chains")
            return chains
            
        except Exception as e:
            self.progress_tracker.stop_tracking(tracking_id, status="failed", message=str(e))
            raise
    
    def _extract_mentions(self, text: str) -> List[Mention]:
        """Extract all mentions from text."""
        mentions = []
        
        # Extract pronouns
        pronoun_patterns = {
            "he": r"\bhe\b",
            "she": r"\bshe\b",
            "it": r"\bit\b",
            "they": r"\bthey\b",
            "his": r"\bhis\b",
            "her": r"\bher\b",
            "their": r"\btheir\b"
        }
        
        for pronoun, pattern in pronoun_patterns.items():
            for match in re.finditer(pattern, text, re.IGNORECASE):
                mentions.append(Mention(
                    text=match.group(0),
                    start_char=match.start(),
                    end_char=match.end(),
                    mention_type="pronoun",
                    metadata={"pronoun": pronoun}
                ))
        
        return mentions
    
    def build_coreference_chains(self, mentions: List[Mention], **options) -> List[CoreferenceChain]:
        """
        Build coreference chains from mentions.
        
        Args:
            mentions: List of mentions
            **options: Chain building options
            
        Returns:
            list: List of coreference chains
        """
        return self.chain_builder.build_coreference_chains(mentions, **options)
    
    def resolve_pronouns(self, text: str, **options) -> List[Tuple[str, str]]:
        """
        Resolve pronoun references in text.
        
        Args:
            text: Input text
            **options: Resolution options
            
        Returns:
            list: List of (pronoun, antecedent) tuples
        """
        mentions = self._extract_mentions(text)
        return self.pronoun_resolver.resolve_pronouns(text, mentions, **options)
    
    def detect_entity_coreferences(self, text: str, entities: List[Entity], **options) -> List[CoreferenceChain]:
        """
        Detect entity coreferences in text.
        
        Args:
            text: Input text
            entities: List of entities
            **options: Detection options
            
        Returns:
            list: List of coreference chains
        """
        # Convert entities to mentions
        mentions = [
            Mention(
                text=e.text,
                start_char=e.start_char,
                end_char=e.end_char,
                mention_type="entity",
                metadata={"entity_label": e.label}
            )
            for e in entities
        ]
        
        return self.entity_detector.detect_entity_coreferences(text, mentions, **options)


class PronounResolver:
    """Pronoun resolution engine."""
    
    def __init__(self, **config):
        """Initialize pronoun resolver."""
        self.logger = get_logger("pronoun_resolver")
        self.config = config
    
    def resolve_pronouns(
        self,
        text: str,
        mentions: List[Mention],
        **options
    ) -> List[Tuple[str, str]]:
        """
        Resolve pronoun references in text.
        
        Args:
            text: Input text
            mentions: List of mentions
            **options: Resolution options
            
        Returns:
            list: List of (pronoun, antecedent) tuples
        """
        resolutions = []
        
        # Get pronouns and entities
        pronouns = [m for m in mentions if m.mention_type == "pronoun"]
        entities = [m for m in mentions if m.mention_type == "entity" or m.mention_type == "nominal"]
        
        # Simple resolution: find closest preceding entity
        for pronoun in pronouns:
            # Find preceding entities
            preceding = [e for e in entities if e.end_char < pronoun.start_char]
            
            if preceding:
                # Take closest
                antecedent = max(preceding, key=lambda e: e.end_char)
                resolutions.append((pronoun.text, antecedent.text))
        
        return resolutions


class EntityCoreferenceDetector:
    """Entity coreference detection."""
    
    def __init__(self, **config):
        """Initialize entity coreference detector."""
        self.logger = get_logger("entity_coreference_detector")
        self.config = config
    
    def detect_entity_coreferences(
        self,
        text: str,
        mentions: List[Mention],
        **options
    ) -> List[CoreferenceChain]:
        """
        Detect entity coreferences in text.
        
        Args:
            text: Input text
            mentions: List of mentions
            **options: Detection options
            
        Returns:
            list: List of coreference chains
        """
        chains = []
        
        # Group mentions by text similarity
        entity_mentions = [m for m in mentions if m.mention_type in ["entity", "nominal"]]
        
        # Simple grouping by exact text match
        text_groups = {}
        for mention in entity_mentions:
            key = mention.text.lower()
            if key not in text_groups:
                text_groups[key] = []
            text_groups[key].append(mention)
        
        # Create chains from groups
        for key, group_mentions in text_groups.items():
            if len(group_mentions) > 1:
                # Use first mention as representative
                representative = group_mentions[0]
                chain = CoreferenceChain(
                    mentions=group_mentions,
                    representative=representative
                )
                chains.append(chain)
        
        return chains


class CoreferenceChainBuilder:
    """Coreference chain construction."""
    
    def __init__(self, **config):
        """Initialize coreference chain builder."""
        self.logger = get_logger("coreference_chain_builder")
        self.config = config
    
    def build_coreference_chains(self, mentions: List[Mention], **options) -> List[CoreferenceChain]:
        """
        Build coreference chains from mentions.
        
        Args:
            mentions: List of mentions
            **options: Chain building options
            
        Returns:
            list: List of coreference chains
        """
        chains = []
        
        # Simple implementation: group by text similarity
        processed = set()
        
        for mention in mentions:
            if mention.text.lower() in processed:
                continue
            
            # Find similar mentions
            similar = [
                m for m in mentions
                if m.text.lower() == mention.text.lower() or
                self._similar_mentions(mention.text, m.text)
            ]
            
            if len(similar) > 1:
                processed.add(mention.text.lower())
                
                # Representative is first (leftmost) mention
                representative = min(similar, key=lambda m: m.start_char)
                
                chain = CoreferenceChain(
                    mentions=similar,
                    representative=representative,
                    entity_type=similar[0].metadata.get("entity_label")
                )
                chains.append(chain)
        
        return chains
    
    def _similar_mentions(self, text1: str, text2: str) -> bool:
        """Check if two mentions are similar."""
        t1_lower = text1.lower()
        t2_lower = text2.lower()
        
        # Exact match
        if t1_lower == t2_lower:
            return True
        
        # One contains the other
        if t1_lower in t2_lower or t2_lower in t1_lower:
            return True
        
        # Word overlap
        words1 = set(t1_lower.split())
        words2 = set(t2_lower.split())
        overlap = len(words1 & words2) / max(len(words1), len(words2)) if words1 or words2 else 0
        
        return overlap > 0.7
